Support async custom completion closures.

Signed-off-by: Ross Goldberg <[email protected]>

Author: rgoldberg

Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift

@@ -380,7 +380,7 @@ extension [ParsableCommand.Type] {
 
         """
 
-    case .custom:
+    case .custom, .customAsync:
       // Generate a call back into the command to retrieve a completions list
       return """
         \(addCompletionsFunctionName) -W\

Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift

@@ -223,7 +223,7 @@ extension [ParsableCommand.Type] {
       results += ["-\(r)fa '(\(completeDirectoriesFunctionName))'"]
     case .shellCommand(let shellCommand):
       results += ["-\(r)fka '(\(shellCommand))'"]
-    case .custom:
+    case .custom, .customAsync:
       results += [
         """
         -\(r)fka '(\

Sources/ArgumentParser/Completions/ZshCompletionsGenerator.swift

@@ -204,7 +204,7 @@ extension [ParsableCommand.Type] {
         nil
       )
 
-    case .custom:
+    case .custom, .customAsync:
       return (
         "{\(customCompleteFunctionName) \(arg.customCompletionCall(self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
         nil

Sources/ArgumentParser/Parsable Properties/CompletionKind.swift

@@ -41,6 +41,7 @@ public struct CompletionKind {
     case directory
     case shellCommand(String)
     case custom(@Sendable ([String], Int, String) -> [String])
+    case customAsync(@Sendable ([String], Int, String) async -> [String])
     case customDeprecated(@Sendable ([String]) -> [String])
   }
 
@@ -176,6 +177,17 @@ public struct CompletionKind {
     CompletionKind(kind: .custom(completion))
   }
 
+  /// Generate completions using the given async closure.
+  ///
+  /// The same as `custom(@Sendable @escaping ([String], Int, String) -> [String])`,
+  /// except that the closure is asynchronous.
+  @available(macOS 10.15, macCatalyst 13, iOS 13, tvOS 13, watchOS 6, *)
+  public static func custom(
+    _ completion: @Sendable @escaping ([String], Int, String) async -> [String]
+  ) -> CompletionKind {
+    CompletionKind(kind: .customAsync(completion))
+  }
+
   /// Deprecated; only kept for backwards compatibility.
   ///
   /// The same as `custom(@Sendable @escaping ([String], Int, String) -> [String])`,

Sources/ArgumentParser/Parsing/CommandParser.swift

@@ -10,8 +10,10 @@
 //===----------------------------------------------------------------------===//
 
 #if swift(>=6.0)
+private import class Dispatch.DispatchSemaphore
 internal import class Foundation.ProcessInfo
 #else
+import class Dispatch.DispatchSemaphore
 import class Foundation.ProcessInfo
 #endif
 
@@ -447,37 +449,20 @@ extension CommandParser {
     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 arg = args.popFirst(),
-        let cursorIndexWithinCompletingArgument = Int(arg)
-      else {
-        throw ParserError.invalidState
-      }
-
-      let completingPrefix: String
-      if let completingArgument = args.last {
-        completingPrefix = String(
-          completingArgument.prefix(cursorIndexWithinCompletingArgument)
-        )
-      } else if cursorIndexWithinCompletingArgument == 0 {
-        completingPrefix = ""
-      } else {
-        throw ParserError.invalidState
-      }
-
+      let (args, completingArgumentIndex, completingPrefix) =
+        try parseCustomCompletionArguments(from: args)
       completions = complete(
-        Array(args),
+        args,
         completingArgumentIndex,
         completingPrefix
       )
+    case .customAsync(let complete):
+      if #available(macOS 10.15, macCatalyst 13, iOS 13, tvOS 13, watchOS 6, *)
+      {
+        completions = try asyncCustomCompletions(from: args, complete: complete)
+      } else {
+        throw ParserError.invalidState
+      }
     case .customDeprecated(let complete):
       completions = complete(args)
     default:
@@ -494,6 +479,57 @@ extension CommandParser {
   }
 }
 
+private func parseCustomCompletionArguments(
+  from args: [String]
+) throws -> ([String], Int, String) {
+  var args = args.dropFirst(0)
+  guard
+    let s = args.popFirst(),
+    let completingArgumentIndex = Int(s)
+  else {
+    throw ParserError.invalidState
+  }
+
+  guard
+    let arg = args.popFirst(),
+    let cursorIndexWithinCompletingArgument = Int(arg)
+  else {
+    throw ParserError.invalidState
+  }
+
+  let completingPrefix: String
+  if let completingArgument = args.last {
+    completingPrefix = String(
+      completingArgument.prefix(cursorIndexWithinCompletingArgument)
+    )
+  } else if cursorIndexWithinCompletingArgument == 0 {
+    completingPrefix = ""
+  } else {
+    throw ParserError.invalidState
+  }
+
+  return (Array(args), completingArgumentIndex, completingPrefix)
+}
+
+@available(macOS 10.15, macCatalyst 13, iOS 13, tvOS 13, watchOS 6, *)
+private func asyncCustomCompletions(
+  from args: [String],
+  complete: @escaping @Sendable ([String], Int, String) async -> [String]
+) throws -> [String] {
+  let (args, completingArgumentIndex, completingPrefix) =
+    try parseCustomCompletionArguments(from: args)
+  var completions: [String] = []
+  let semaphore = DispatchSemaphore(value: 0)
+  Task {
+    completions = await complete(
+      args, completingArgumentIndex, completingPrefix
+    )
+    semaphore.signal()
+  }
+  semaphore.wait()
+  return completions
+}
+
 // MARK: Building Command Stacks
 
 extension CommandParser {

Sources/ArgumentParser/Usage/DumpHelpGenerator.swift

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

Sources/ArgumentParserToolInfo/ToolInfo.swift

@@ -151,6 +151,8 @@ public struct ArgumentInfoV0: Codable, Hashable {
     case shellCommand(command: String)
     /// Generate completions using the given three-parameter closure.
     case custom
+    /// Generate completions using the given async three-parameter closure.
+    case customAsync
     /// Generate completions using the given one-parameter closure.
     @available(*, deprecated, message: "Use custom instead.")
     case customDeprecated