Repository: apple/swift-argument-parser (swift)

Issue:
Zsh completion missing for arguments in an option group with custom completion
Zsh completions are not shown for arguments that are added to a command using an `@OptionGroup` and that have a `custom` completion kind.

**ArgumentParser version:** `main`
**Swift version:** 
```
swift-driver version: 1.90.11.1 Apple Swift version 5.10 (swiftlang-5.10.0.13 clang-1500.3.9.4)
Target: arm64-apple-macosx14.0
```

### Checklist
- [x] If possible, I've reproduced the issue using the `main` branch of this package
- [x] I've searched for [existing GitHub issues](https://github.com/apple/swift-argument-parser/issues)

### Steps to Reproduce

1. Create a command with an option group with an arguments with a custom completion kind.
	```swift
	import ArgumentParser
	
	struct ArgOptionGroup: ParsableArguments {
	    @Argument(completion: .custom { _ in
	        ["1-argInOptionGroup", "2-argInOptionGroup", "3-argInOptionGroup"]
	    })
	    var argInOptionGroup: String
	}
	
	@main
	struct CompletionExample: ParsableCommand {
	    @Argument(completion: .custom { _ in
	        ["1-argInCommand", "2-argInCommand", "3-argInCommand"]
	    })
	    var argInCommand: String
	    
	    @OptionGroup var optionGroup: ArgOptionGroup
	}
	```

1. Install the completions in your terminal and load them, steps will vary per environment, for me the steps are.
	1. `swift build` 
	1. Install the completions.
		```shell
		./completion-example --generate-completion-script zsh > ~/.oh-my-zsh/completions/_completion-example
		```
    1. Open a new terminal to load the new completions.
1. Type `./completion-example ` (note the trailing space)
	<img width='50%' src='https://github.com/apple/swift-argument-parser/assets/680442/32163054-d25a-41ef-ad0a-7188e76e6591'>

1. Press tab to accept the first completion
1. Press space.
	<img width='50%' src='https://github.com/apple/swift-argument-parser/assets/680442/d954b235-8a71-4bcb-876f-eeffffab3f7d'>
1. Type `1` (the first character of one of the arguments.
	<img width='50%' src='https://github.com/apple/swift-argument-parser/assets/680442/82ebd942-f5af-472e-9f56-fd88aca67bd8'>

### Expected behavior
It should suggest completions for the argument from the option group, as it did for the argument in the command (step 3).

### Actual behavior
It doesn't show any completions for the arguments in the option group (step 4). When starting to type any of the options that should be suggested (step 5), it prints that there are no matches (probably from `zsh-autosuggestions`).

Pull request:
## Issue
- Fixes #646 

## Description

### The issue

It was caused by a mismatch in the completion scripts and the custom completion parsing. This is following the example from the issue above.
- The completion scripts contained the name of the argument
  ```zsh
  ':arg-in-option-group:{_custom_completion $_completion_example_commandname ---completion  -- argInOptionGroup $words}'
  ```
- The parsing of the completion command failed to find the argument using `ArgumentSet.firstPositional(named:)`.
  This was caused by the function creating an `InputKey` with just a `name` and comparing to keys in the set that have their `name` and `path` set.

  In other words
  ```swift
  InputKey(name: "argInOptionGroup", parent: nil)
  != InputKey(name: "argInOptionGroup", parent: .init(named: "optionGroup", parent: nil))
  ```

There's also another issue that would occur when a command and an option group each contain an argument with the same name. In that case the completion would always match the first one in the `ArgumentSet` even if the user is trying to get completions for the second argument.

## The fix

The completions scripts were updated to contain the full path of the argument instead of just the name. Then the completions parser can create an `InputKey` that will exactly match one in the `ArgumentSet`.

To avoid having the logic spread out, an extension was added to `InputKey` to get the `fullPathString` from a key and an init to create an `InputKey` from a `fullPathString`.

I also updated the bash completion script to use `ArgumentDefinition.customCompletionCall(_:)` that the zsh script was already using to deduplicate the logic.

### Checklist
- [x] I've added at least one test that validates that my change is working, if appropriate
- [x] I've followed the code style of the rest of the project
- [x] I've read the [Contribution Guidelines](https://github.com/apple/swift-argument-parser/blob/main/CONTRIBUTING.md)
- [x] I've updated the documentation if necessary (wasn't necessary)

Patch:
diff --git a/Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift b/Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift
index 2bebc19ec..cd7864b5c 100644
--- a/Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift
+++ b/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[@]}")
             """
         }
       }
diff --git a/Sources/ArgumentParser/Completions/CompletionsGenerator.swift b/Sources/ArgumentParser/Completions/CompletionsGenerator.swift
index d411cff4a..bbfe0bb72 100644
--- a/Sources/ArgumentParser/Completions/CompletionsGenerator.swift
+++ b/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)"
   }
 }
diff --git a/Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift b/Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift
index 9a8fa83dd..454ca2a9a 100644
--- a/Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift
+++ b/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,12 +45,12 @@ 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
@@ -58,37 +58,40 @@ extension FishCompletionsGenerator {
 }
 
 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
   }
 }
 
diff --git a/Sources/ArgumentParser/Parsing/ArgumentSet.swift b/Sources/ArgumentParser/Parsing/ArgumentSet.swift
index 9490f1e63..1289245ca 100644
--- a/Sources/ArgumentParser/Parsing/ArgumentSet.swift
+++ b/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) })
   }
 }
 
diff --git a/Sources/ArgumentParser/Parsing/CommandParser.swift b/Sources/ArgumentParser/Parsing/CommandParser.swift
index 0fa83f6ed..c09220b34 100644
--- a/Sources/ArgumentParser/Parsing/CommandParser.swift
+++ b/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
diff --git a/Sources/ArgumentParser/Parsing/InputKey.swift b/Sources/ArgumentParser/Parsing/InputKey.swift
index 9837e41e0..be71a11e6 100644
--- a/Sources/ArgumentParser/Parsing/InputKey.swift
+++ b/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()
   }
 }


Test patch:
diff --git a/Tests/ArgumentParserExampleTests/MathExampleTests.swift b/Tests/ArgumentParserExampleTests/MathExampleTests.swift
index 0a9c09a38..ee677a6a0 100644
--- a/Tests/ArgumentParserExampleTests/MathExampleTests.swift
+++ b/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.\'
diff --git a/Tests/ArgumentParserUnitTests/CompletionScriptTests.swift b/Tests/ArgumentParserUnitTests/CompletionScriptTests.swift
index 5a8d8dfe7..0c17a2965 100644
--- a/Tests/ArgumentParserUnitTests/CompletionScriptTests.swift
+++ b/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.'