Fix zsh/bash completions for arguments in option groups with a custom completion

Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift

@@ -144,14 +144,8 @@ struct BashCompletionsGenerator {
         case .shellCommand(let command):
           return "$(\(command))"
         case .custom:
-          // Generate a call back into the command to retrieve a completions list
-          let subcommandNames = commands.dropFirst().map { $0._commandName }.joined(separator: " ")
-          // TODO: Make this work for @Arguments
-          let argumentName = arg.names.preferredName?.synopsisString
-                ?? arg.help.keys.first?.name ?? "---"
-
           return """
-            $("${COMP_WORDS[0]}" ---completion \(subcommandNames) -- \(argumentName) "${COMP_WORDS[@]}")
+            $("${COMP_WORDS[0]}" \(arg.customCompletionCall(commands)) "${COMP_WORDS[@]}")
             """
         }
       }

Sources/ArgumentParser/Completions/CompletionsGenerator.swift

@@ -88,7 +88,7 @@ extension ArgumentDefinition {
   func customCompletionCall(_ commands: [ParsableCommand.Type]) -> String {
     let subcommandNames = commands.dropFirst().map { $0._commandName }.joined(separator: " ")
     let argumentName = names.preferredName?.synopsisString
-          ?? self.help.keys.first?.name ?? "---"
+          ?? self.help.keys.first?.fullPathString ?? "---"
     return "---completion \(subcommandNames) -- \(argumentName)"
   }
 }

Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift

@@ -5,7 +5,7 @@ struct FishCompletionsGenerator {
       preprocessorFunction(commandName: programName),
       helperFunction(commandName: programName)
     ]
-    let completions = generateCompletions(commandChain: [programName], [type])
+    let completions = generateCompletions([type])
 
     return helperFunctions.joined(separator: "\n\n") + "\n\n" + completions.joined(separator: "\n")
   }
@@ -14,10 +14,10 @@ struct FishCompletionsGenerator {
 // MARK: - Private functions
 
 extension FishCompletionsGenerator {
-  private static func generateCompletions(commandChain: [String], _ commands: [ParsableCommand.Type]) -> [String] {
+  private static func generateCompletions(_ commands: [ParsableCommand.Type]) -> [String] {
     let type = commands.last!
     let isRootCommand = commands.count == 1
-    let programName = commandChain[0]
+    let programName = commands[0]._commandName
     var subcommands = type.configuration.subcommands
       .filter { $0.configuration.shouldDisplay }
 
@@ -27,7 +27,7 @@ extension FishCompletionsGenerator {
 
     let helperFunctionName = helperFunctionName(commandName: programName)
 
-    var prefix = "complete -c \(programName) -n '\(helperFunctionName) \"\(commandChain.joined(separator: separator))\""
+    var prefix = "complete -c \(programName) -n '\(helperFunctionName) \"\(commands.map { $0._commandName }.joined(separator: separator))\""
     if !subcommands.isEmpty {
       prefix += " \"\(subcommands.map { $0._commandName }.joined(separator: separator))\""
     }
@@ -45,50 +45,53 @@ extension FishCompletionsGenerator {
 
     let argumentCompletions = commands
       .argumentsForHelp(visibility: .default)
-      .compactMap { $0.argumentSegments(commandChain) }
+      .compactMap { $0.argumentSegments(commands) }
       .map { $0.joined(separator: " ") }
       .map { complete(suggestion: $0) }
 
     let completionsFromSubcommands = subcommands.flatMap { subcommand in
-      generateCompletions(commandChain: commandChain + [subcommand._commandName], [subcommand])
+      generateCompletions(commands + [subcommand])
     }
 
     return completionsFromSubcommands + argumentCompletions + subcommandCompletions
   }
 }
 
 extension ArgumentDefinition {
-  fileprivate func argumentSegments(_ commandChain: [String]) -> [String]? {
-    guard help.visibility.base == .default,
-          !names.isEmpty
+  fileprivate func argumentSegments(_ commands: [ParsableCommand.Type]) -> [String]? {
+    guard help.visibility.base == .default
     else { return nil }
-
-    var results = names.map{ $0.asFishSuggestion }
-
+
+    var results: [String] = []
+
+    if !names.isEmpty {
+      results += names.map{ $0.asFishSuggestion }
+    }
+
     if !help.abstract.isEmpty {
       results += ["-d '\(help.abstract.fishEscape())'"]
     }
-
-    if isNullary {
-      return results
-    }
-
+
     switch completion.kind {
-    case .default: return results
+    case .default where names.isEmpty:
+      return nil
+    case .default:
+      break
     case .list(let list):
-      return results + ["-r -f -k -a '\(list.joined(separator: " "))'"]
+      results += ["-r -f -k -a '\(list.joined(separator: " "))'"]
     case .file(let extensions):
       let pattern = "*.{\(extensions.joined(separator: ","))}"
-      return results + ["-r -f -a '(for i in \(pattern); echo $i;end)'"]
+      results += ["-r -f -a '(for i in \(pattern); echo $i;end)'"]
     case .directory:
-      return results + ["-r -f -a '(__fish_complete_directories)'"]
+      results += ["-r -f -a '(__fish_complete_directories)'"]
     case .shellCommand(let shellCommand):
-      return results + ["-r -f -a '(\(shellCommand))'"]
+      results += ["-r -f -a '(\(shellCommand))'"]
     case .custom:
-      let program = commandChain[0]
-      let subcommands = commandChain.dropFirst().joined(separator: " ")
-      return results + ["-r -f -a '(command \(program) ---completion \(subcommands) -- --custom (commandline -opc)[1..-1])'"]
+      let commandName = commands.first!._commandName
+      results += ["-r -f -a '(command \(commandName) \(customCompletionCall(commands)) (commandline -opc)[1..-1])'"]
     }
+
+    return results
   }
 }
 

Sources/ArgumentParser/Parsing/ArgumentSet.swift

@@ -191,10 +191,9 @@ extension ArgumentSet {
   }
 
   func firstPositional(
-    named name: String
+    withKey key: InputKey
   ) -> ArgumentDefinition? {
-    let key = InputKey(name: name, parent: nil)
-    return first(where: { $0.help.keys.contains(key) })
+    first(where: { $0.help.keys.contains(key) })
   }
 }
 

Sources/ArgumentParser/Parsing/CommandParser.swift

@@ -343,7 +343,8 @@ extension CommandParser {
       completionFunction = f
 
     case .value(let str):
-      guard let matchedArgument = argset.firstPositional(named: str),
+      guard let key = InputKey(fullPathString: str),
+        let matchedArgument = argset.firstPositional(withKey: key),
         case .custom(let f) = matchedArgument.completion.kind
         else { throw ParserError.invalidState }
       completionFunction = f

Sources/ArgumentParser/Parsing/InputKey.swift

@@ -51,6 +51,23 @@ struct InputKey: Hashable {
 
 extension InputKey: CustomStringConvertible {
   var description: String {
-    fullPath.joined(separator: ".")
+      fullPathString
+  }
+}
+
+extension InputKey {
+  private static var separator: Character { "." }
+
+  var fullPathString: String {
+    fullPath.joined(separator: .init(Self.separator))
+  }
+
+  init?(fullPathString: String) {
+    let fullPath = fullPathString.split(separator: Self.separator).map(String.init)
+
+    guard let name = fullPath.last else { return nil }
+
+    self.name = name
+    self.path = fullPath.dropLast()
   }
 }

Tests/ArgumentParserExampleTests/MathExampleTests.swift

@@ -608,23 +608,30 @@ function _swift_math_using_command
 end
 
 complete -c math -n \'_swift_math_using_command \"math add\"\' -l hex-output -s x -d \'Use hexadecimal notation for the result.\'
+complete -c math -n \'_swift_math_using_command \"math add\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math add\"\' -s h -l help -d \'Show help information.\'
 complete -c math -n \'_swift_math_using_command \"math multiply\"\' -l hex-output -s x -d \'Use hexadecimal notation for the result.\'
+complete -c math -n \'_swift_math_using_command \"math multiply\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math multiply\"\' -s h -l help -d \'Show help information.\'
 complete -c math -n \'_swift_math_using_command \"math stats average\"\' -l kind -d \'The kind of average to provide.\' -r -f -k -a \'mean median mode\'
 complete -c math -n \'_swift_math_using_command \"math stats average\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math stats average\"\' -s h -l help -d \'Show help information.\'
+complete -c math -n \'_swift_math_using_command \"math stats stdev\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math stats stdev\"\' -s h -l help -d \'Show help information.\'
+complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -r -f -k -a \'alphabet alligator branch braggart\'
+complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -r -f -a \'(command math ---completion stats quantiles -- customArg (commandline -opc)[1..-1])\'
 complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -l file -r -f -a \'(for i in *.{txt,md}; echo $i;end)\'
 complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -l directory -r -f -a \'(__fish_complete_directories)\'
 complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -l shell -r -f -a \'(head -100 /usr/share/dict/words | tail -50)\'
 complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -l custom -r -f -a \'(command math ---completion stats quantiles -- --custom (commandline -opc)[1..-1])\'
+complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math stats quantiles\"\' -s h -l help -d \'Show help information.\'
-complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles help\"\' -s h -l help -d \'Show help information.\'
-complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles help\"\' -f -a \'average\' -d \'Print the average of the values.\'
-complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles help\"\' -f -a \'stdev\' -d \'Print the standard deviation of the values.\'
-complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles help\"\' -f -a \'quantiles\' -d \'Print the quantiles of the values (TBD).\'
-complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles help\"\' -f -a \'help\' -d \'Show subcommand help information.\'
+complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles\"\' -l version -d \'Show the version.\'
+complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles\"\' -s h -l help -d \'Show help information.\'
+complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles\"\' -f -a \'average\' -d \'Print the average of the values.\'
+complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles\"\' -f -a \'stdev\' -d \'Print the standard deviation of the values.\'
+complete -c math -n \'_swift_math_using_command \"math stats\" \"average stdev quantiles\"\' -f -a \'quantiles\' -d \'Print the quantiles of the values (TBD).\'
+complete -c math -n \'_swift_math_using_command \"math help\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math\" \"add multiply stats help\"\' -l version -d \'Show the version.\'
 complete -c math -n \'_swift_math_using_command \"math\" \"add multiply stats help\"\' -s h -l help -d \'Show help information.\'
 complete -c math -n \'_swift_math_using_command \"math\" \"add multiply stats help\"\' -f -a \'add\' -d \'Print the sum of the values.\'

Tests/ArgumentParserUnitTests/CompletionScriptTests.swift

@@ -33,6 +33,11 @@ extension CompletionScriptTests {
     case one, two, three = "custom-three"
   }
 
+  struct NestedArguments: ParsableArguments {
+    @Argument(completion: .custom { _ in ["t", "u", "v"] })
+    var nestedArgument: String
+  }
+
   struct Base: ParsableCommand {
     static let configuration = CommandConfiguration(
       commandName: "base-test",
@@ -53,12 +58,15 @@ extension CompletionScriptTests {
 
     @Option() var rep1: [String]
     @Option(name: [.short, .long]) var rep2: [String]
-
-   struct SubCommand: ParsableCommand {
-     static let configuration = CommandConfiguration(
-       commandName: "sub-command"
-     )
-   }
+
+    @Argument(completion: .custom { _ in ["q", "r", "s"] }) var argument: String
+    @OptionGroup var nested: NestedArguments
+
+    struct SubCommand: ParsableCommand {
+      static let configuration = CommandConfiguration(
+        commandName: "sub-command"
+      )
+    }
   }
 
   func testBase_Zsh() throws {
@@ -111,6 +119,13 @@ extension CompletionScriptTests {
 
     @Option(name: .customShort("z"), completion: .custom { _ in ["x", "y", "z"] })
     var three: String
+
+    @OptionGroup var nested: NestedArguments
+
+    struct NestedArguments: ParsableArguments {
+      @Argument(completion: .custom { _ in ["g", "h", "i"] })
+      var four: String
+    }
   }
 
   func verifyCustomOutput(
@@ -134,8 +149,10 @@ extension CompletionScriptTests {
     try verifyCustomOutput("--one", expectedOutput: "a\nb\nc")
     try verifyCustomOutput("two", expectedOutput: "d\ne\nf")
     try verifyCustomOutput("-z", expectedOutput: "x\ny\nz")
+    try verifyCustomOutput("nested.four", expectedOutput: "g\nh\ni")
 
     XCTAssertThrowsError(try verifyCustomOutput("--bad", expectedOutput: ""))
+    XCTAssertThrowsError(try verifyCustomOutput("four", expectedOutput: ""))
   }
 }
 
@@ -175,6 +192,8 @@ _base-test() {
         '*--kind-counter'
         '*--rep1:rep1:'
         '*'{-r,--rep2}':rep2:'
+        ':argument:{_custom_completion $_base_test_commandname ---completion  -- argument $words}'
+        ':nested-argument:{_custom_completion $_base_test_commandname ---completion  -- nested.nestedArgument $words}'
         '(-h --help)'{-h,--help}'[Show help information.]'
         '(-): :->command'
         '(-)*:: :->arg'
@@ -243,6 +262,8 @@ _base_test() {
     prev="${COMP_WORDS[COMP_CWORD-1]}"
     COMPREPLY=()
     opts="--name --kind --other-kind --path1 --path2 --path3 --one --two --three --kind-counter --rep1 -r --rep2 -h --help sub-command help"
+    opts="$opts $("${COMP_WORDS[0]}" ---completion  -- argument "${COMP_WORDS[@]}")"
+    opts="$opts $("${COMP_WORDS[0]}" ---completion  -- nested.nestedArgument "${COMP_WORDS[@]}")"
     if [[ $COMP_CWORD == "1" ]]; then
         COMPREPLY=( $(compgen -W "$opts" -- "$cur") )
         return
@@ -401,6 +422,8 @@ complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-comman
 complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -l kind-counter
 complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -l rep1
 complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -s r -l rep2
+complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -r -f -a '(command base-test ---completion  -- argument (commandline -opc)[1..-1])'
+complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -r -f -a '(command base-test ---completion  -- nested.nestedArgument (commandline -opc)[1..-1])'
 complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -s h -l help -d 'Show help information.'
 complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -f -a 'sub-command' -d ''
 complete -c base-test -n '_swift_base-test_using_command "base-test" "sub-command help"' -f -a 'help' -d 'Show subcommand help information.'