Add 2 index arguments to custom shell completion calls (#763)

They indicate to the Swift custom completion function:

1. the word for which completions are being requested.
2. the location of the cursor within that word.

Author: rgoldberg

Examples/math/Math.swift

@@ -198,10 +198,17 @@ extension Math.Statistics {
     var oneOfFour: String?
 
     @Argument(
-      completion: .custom { _ in ["alabaster", "breakfast", "crunch", "crash"] }
+      completion: .custom { _, _, _ in
+        ["alabaster", "breakfast", "crunch", "crash"]
+      }
     )
     var customArg: String?
 
+    @Argument(
+      completion: .custom { _ in ["alabaster", "breakfast", "crunch", "crash"] }
+    )
+    var customDeprecatedArg: String?
+
     @Argument(help: "A group of floating-point values to operate on.")
     var values: [Double] = []
 
@@ -222,12 +229,16 @@ extension Math.Statistics {
     var directory: String?
 
     @Option(
-      completion: .shellCommand("head -100 /usr/share/dict/words | tail -50"))
+      completion: .shellCommand("head -100 /usr/share/dict/words | tail -50")
+    )
     var shell: String?
 
     @Option(completion: .custom(customCompletion))
     var custom: String?
 
+    @Option(completion: .custom(customDeprecatedCompletion))
+    var customDeprecated: String?
+
     func validate() throws {
       if testSuccessExitCode {
         throw ExitCode.success
@@ -248,7 +259,13 @@ extension Math.Statistics {
   }
 }
 
-func customCompletion(_ s: [String]) -> [String] {
+func customCompletion(_ s: [String], _: Int, _: Int) -> [String] {
+  (s.last ?? "").starts(with: "a")
+    ? ["aardvark", "aaaaalbert"]
+    : ["hello", "helicopter", "heliotrope"]
+}
+
+func customDeprecatedCompletion(_ s: [String]) -> [String] {
   (s.last ?? "").starts(with: "a")
     ? ["aardvark", "aaaaalbert"]
     : ["hello", "helicopter", "heliotrope"]

Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift

@@ -21,6 +21,22 @@ extension [ParsableCommand.Type] {
     return """
       #!/bin/bash
 
+      \(cursorIndexInCurrentWordFunctionName)() {
+          local remaining="${COMP_LINE}"
+
+          local word
+          for word in "${COMP_WORDS[@]::COMP_CWORD}"; do
+              remaining="${remaining##*([[:space:]])"${word}"*([[:space:]])}"
+          done
+
+          local -ir index="$((COMP_POINT - ${#COMP_LINE} + ${#remaining}))"
+          if [[ "${index}" -le 0 ]]; then
+              printf 0
+          else
+              printf %s "${index}"
+          fi
+      }
+
       # positional arguments:
       #
       # - 1: the current (sub)command's count of positional arguments
@@ -365,6 +381,16 @@ extension [ParsableCommand.Type] {
         """
 
     case .custom:
+      // Generate a call back into the command to retrieve a completions list
+      return """
+        \(addCompletionsFunctionName) -W\
+         "$(\(customCompleteFunctionName) \(arg.customCompletionCall(self))\
+         "${COMP_CWORD}"\
+         "$(\(cursorIndexInCurrentWordFunctionName))")"
+
+        """
+
+    case .customDeprecated:
       // Generate a call back into the command to retrieve a completions list
       return """
         \(addCompletionsFunctionName) -W\
@@ -374,6 +400,10 @@ extension [ParsableCommand.Type] {
     }
   }
 
+  private var cursorIndexInCurrentWordFunctionName: String {
+    "_\(prefix(1).completionFunctionName().shellEscapeForVariableName())_cursor_index_in_current_word"
+  }
+
   private var offerFlagsOptionsFunctionName: String {
     "_\(prefix(1).completionFunctionName().shellEscapeForVariableName())_offer_flags_options"
   }

Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift

@@ -224,6 +224,15 @@ extension [ParsableCommand.Type] {
     case .shellCommand(let shellCommand):
       results += ["-\(r)fka '(\(shellCommand))'"]
     case .custom:
+      results += [
+        """
+        -\(r)fka '(\
+        \(customCompletionFunctionName) \(arg.customCompletionCall(self)) \
+        (count (\(tokensFunctionName) -pc)) (\(tokensFunctionName) -tC)\
+        )'
+        """
+      ]
+    case .customDeprecated:
       results += [
         """
         -\(r)fka '(\(customCompletionFunctionName) \(arg.customCompletionCall(self)))'

Sources/ArgumentParser/Completions/ZshCompletionsGenerator.swift

@@ -28,12 +28,20 @@ extension [ParsableCommand.Type] {
 
       \(customCompleteFunctionName)() {
           local -a completions
-          completions=("${(@f)"$("${@}")"}")
+          completions=("${(@f)"$("${command_name}" "${@}" "${command_line[@]}")"}")
           if [[ "${#completions[@]}" -gt 1 ]]; then
               \(completeFunctionName) "${completions[@]:0:-1}"
           fi
       }
 
+      \(cursorIndexInCurrentWordFunctionName)() {
+          if [[ -z "${QIPREFIX}${IPREFIX}${PREFIX}" ]]; then
+              printf 0
+          else
+              printf %s "${#${(z)LBUFFER}[-1]}"
+          fi
+      }
+
       \(completionFunctions)\
       \(completionFunctionName())
       """
@@ -107,6 +115,7 @@ extension [ParsableCommand.Type] {
 
               local -r command_name="${words[1]}"
               local -ar command_line=("${words[@]}")
+              local -ir current_word_index="$((CURRENT - 1))"
 
 
           """
@@ -192,7 +201,13 @@ extension [ParsableCommand.Type] {
 
     case .custom:
       return (
-        "{\(customCompleteFunctionName) \"${command_name}\" \(arg.customCompletionCall(self)) \"${command_line[@]}\"}",
+        "{\(customCompleteFunctionName) \(arg.customCompletionCall(self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
+        nil
+      )
+
+    case .customDeprecated:
+      return (
+        "{\(customCompleteFunctionName) \(arg.customCompletionCall(self))}",
         nil
       )
     }
@@ -218,6 +233,10 @@ extension [ParsableCommand.Type] {
     // Precondition: first is guaranteed to be non-empty
     "__\(first!._commandName)_custom_complete"
   }
+
+  private var cursorIndexInCurrentWordFunctionName: String {
+    "__\(first?._commandName ?? "")_cursor_index_in_current_word"
+  }
 }
 
 extension String {

Sources/ArgumentParser/Parsable Properties/CompletionKind.swift

@@ -40,7 +40,8 @@ public struct CompletionKind {
     case file(extensions: [String])
     case directory
     case shellCommand(String)
-    case custom(@Sendable ([String]) -> [String])
+    case custom(@Sendable ([String], Int, Int) -> [String])
+    case customDeprecated(@Sendable ([String]) -> [String])
   }
 
   internal var kind: Kind
@@ -125,6 +126,12 @@ public struct CompletionKind {
   /// passed to Swift as `"abc\\""def"` (i.e. the Swift String's contents would
   /// include all 4 of the double quotes and the 2 consecutive backslashes).
   ///
+  /// The first of the two `Int` arguments is the 0-based index of the word
+  /// for which completions are being requested within the given `[String]`.
+  ///
+  /// The second of the two `Int` arguments is the 0-based index of the shell
+  /// cursor within the word for which completions are being requested.
+  ///
   /// ### bash
   ///
   /// In bash 3-, a process substitution (`<(…)`) in the command line prevents
@@ -151,15 +158,40 @@ public struct CompletionKind {
   /// example, the shell word `"abc\\""def"` would be passed to Swift as
   /// `abc\def`. This is fixed in fish 4+.
   ///
+  /// In fish 3-, the cursor index is provided based on the verbatim word, not
+  /// based on the unquoted word, so it can be inconsistent with the unquoted
+  /// word that is supplied to Swift. This problem does not exist in fish 4+.
+  ///
   /// ### zsh
   ///
   /// In zsh, redirects (both their symbol and source/target) are omitted.
+  ///
+  /// In zsh, if the cursor is between a backslash and the character that it
+  /// escapes, the shell cursor index will be indicated as after the escaped
+  /// character, not as after the backslash.
   @preconcurrency
   public static func custom(
-    _ completion: @Sendable @escaping ([String]) -> [String]
+    _ completion: @Sendable @escaping ([String], Int, Int) -> [String]
   ) -> CompletionKind {
     CompletionKind(kind: .custom(completion))
   }
+
+  /// Deprecated; only kept for backwards compatibility.
+  ///
+  /// The same as `custom(@Sendable @escaping ([String], Int, Int) -> [String])`,
+  /// except that index arguments are not supplied.
+  @preconcurrency
+  @available(
+    *,
+    deprecated,
+    message:
+      "Provide a three-parameter closure instead. See custom(@Sendable @escaping ([String], Int, Int) -> [String])."
+  )
+  public static func custom(
+    _ completion: @Sendable @escaping ([String]) -> [String]
+  ) -> CompletionKind {
+    CompletionKind(kind: .customDeprecated(completion))
+  }
 }
 
 extension CompletionKind: Sendable {}

Sources/ArgumentParser/Parsing/CommandParser.swift

@@ -373,7 +373,12 @@ extension CommandParser {
   func handleCustomCompletion(_ arguments: [String]) throws {
     // Completion functions use a custom format:
     //
-    // <command> ---completion [<subcommand> ...] -- <argument-name> [<completion-text>]
+    // <command> ---completion [<subcommand> ...] -- <argument-name> <argument-index> <cursor-index> [<argument> ...]
+    //
+    // <argument-index> is the 0-based index of the <argument> for which completions are being requested.
+    //
+    // <cursor-index> is the 0-based index of the character within the <argument> before which the cursor is located.
+    // For an <argument> whose length is n, if the cursor is after the last element, <cursor-index> will be set to n.
     //
     // The triple-dash prefix makes '---completion' invalid syntax for regular
     // arguments, so it's safe to use for this internal purpose.
@@ -395,35 +400,38 @@ extension CommandParser {
     guard let argToMatch = args.popFirst() else {
       throw ParserError.invalidState
     }
-    // Completion text is optional here
-    let completionValues = Array(args)
 
     // Generate the argument set and parse the argument to find in the set
     let argset = ArgumentSet(current.element, visibility: .private, parent: nil)
     guard let parsedArgument = try parseIndividualArg(argToMatch, at: 0).first
     else { throw ParserError.invalidState }
 
-    // Look up the specified argument and retrieve its custom completion function
-    let completionFunction: ([String]) -> [String]
-
+    // Look up the specified argument, then retrieve & run its custom completion function
     switch parsedArgument.value {
     case .option(let parsed):
-      guard let matchedArgument = argset.first(matching: parsed),
-        case .custom(let f) = matchedArgument.completion.kind
-      else { throw ParserError.invalidState }
-      completionFunction = f
+      guard let matchedArgument = argset.first(matching: parsed) else {
+        throw ParserError.invalidState
+      }
+      try customComplete(matchedArgument, forArguments: Array(args))
 
     case .value(let 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
+      guard
+        let key = InputKey(fullPathString: str),
+        let matchedArgument = argset.firstPositional(withKey: key)
+      else {
+        throw ParserError.invalidState
+      }
+      try customComplete(matchedArgument, forArguments: Array(args))
 
     case .terminator:
       throw ParserError.invalidState
     }
+  }
 
+  private func customComplete(
+    _ argument: ArgumentDefinition,
+    forArguments args: [String]
+  ) throws {
     let environment = ProcessInfo.processInfo.environment
     if let completionShellName = environment[
       CompletionShell.shellEnvironmentVariableName]
@@ -436,10 +444,38 @@ extension CommandParser {
       $0 = environment[CompletionShell.shellVersionEnvironmentVariableName]
     }
 
+    let completions: [String]
+    switch argument.completion.kind {
+    case .custom(let complete):
+      var args = args.dropFirst(0)
+      guard
+        let s = args.popFirst(),
+        let completingArgumentIndex = Int(s)
+      else {
+        throw ParserError.invalidState
+      }
+
+      guard
+        let s = args.popFirst(),
+        let cursorIndexWithinCompletingArgument = Int(s)
+      else {
+        throw ParserError.invalidState
+      }
+
+      completions = complete(
+        Array(args),
+        completingArgumentIndex,
+        cursorIndexWithinCompletingArgument
+      )
+    case .customDeprecated(let complete):
+      completions = complete(args)
+    default:
+      throw ParserError.invalidState
+    }
+
     // Parsing and retrieval successful! We don't want to continue with any
     // other parsing here, so after printing the result of the completion
     // function, exit with a success code.
-    let completions = completionFunction(completionValues)
     throw ParserError.completionScriptCustomResponse(
       CompletionShell.requesting?.format(completions: completions)
         ?? completions.joined(separator: "\n")
@@ -472,9 +508,9 @@ extension CommandParser {
     return result
   }
 
-  func commandStack(for subcommand: ParsableCommand.Type)
-    -> [ParsableCommand.Type]
-  {
+  func commandStack(
+    for subcommand: ParsableCommand.Type
+  ) -> [ParsableCommand.Type] {
     let path = commandTree.path(to: subcommand)
     return path.isEmpty
       ? [commandTree.element]

Sources/ArgumentParser/Usage/DumpHelpGenerator.swift

@@ -224,6 +224,8 @@ extension ArgumentInfoV0.CompletionKindV0 {
       self = .shellCommand(command: command)
     case .custom(_):
       self = .custom
+    case .customDeprecated(_):
+      self = .customDeprecated
     }
   }
 }

Sources/ArgumentParserToolInfo/ToolInfo.swift

@@ -149,8 +149,11 @@ public struct ArgumentInfoV0: Codable, Hashable {
     case directory
     /// Call the given shell command to generate completions.
     case shellCommand(command: String)
-    /// Generate completions using the given closure.
+    /// Generate completions using the given closure including index arguments.
     case custom
+    /// Generate completions using the given closure without index arguments.
+    @available(*, deprecated, message: "Use custom instead.")
+    case customDeprecated
   }
 
   /// Kind of argument the ArgumentInfo describes.