Support async custom completion closures (#782)

Signed-off-by: Ross Goldberg <484615+rgoldberg@users.noreply.github.com>

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,12 @@
 //===----------------------------------------------------------------------===//
 
 #if swift(>=6.0)
+@preconcurrency private import class Dispatch.DispatchSemaphore
+internal import class Foundation.NSLock
 internal import class Foundation.ProcessInfo
 #else
+@preconcurrency import class Dispatch.DispatchSemaphore
+import class Foundation.NSLock
 import class Foundation.ProcessInfo
 #endif
 
@@ -447,37 +451,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 +481,85 @@ 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)
+
+  let completionsBox = SendableBox<[String]>([])
+  let semaphore = DispatchSemaphore(value: 0)
+
+  Task {
+    completionsBox.value = await complete(
+      args,
+      completingArgumentIndex,
+      completingPrefix
+    )
+    semaphore.signal()
+  }
+
+  semaphore.wait()
+  return completionsBox.value
+}
+
+// Helper class to make values sendable across concurrency boundaries
+private final class SendableBox<T>: @unchecked Sendable {
+  private let lock = NSLock()
+  private var _value: T
+
+  init(_ value: T) {
+    self._value = value
+  }
+
+  var value: T {
+    get {
+      lock.lock()
+      defer { lock.unlock() }
+      return _value
+    }
+    set {
+      lock.lock()
+      defer { lock.unlock() }
+      _value = newValue
+    }
+  }
+}
+
 // 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

Tests/ArgumentParserUnitTests/CompletionScriptTests.swift

@@ -27,6 +27,10 @@ private func candidates(prefix: String) -> [String] {
   }
 }
 
+private func candidatesAsync(prefix: String) async -> [String] {
+  candidates(prefix: prefix)
+}
+
 final class CompletionScriptTests: XCTestCase {}
 
 // swift-format-ignore: AlwaysUseLowerCamelCase
@@ -168,6 +172,11 @@ extension CompletionScriptTests {
       @Argument(completion: .custom { _, _, _ in candidates(prefix: "h") })
       var four: String
     }
+
+    @Argument(
+      completion: .custom { _, _, _ in await candidatesAsync(prefix: "j") }
+    )
+    var five: String
   }
 
   func assertCustomCompletion(
@@ -217,6 +226,8 @@ extension CompletionScriptTests {
       "-z", shell: shell, prefix: "g", file: file, line: line)
     try assertCustomCompletion(
       "nested.four", shell: shell, prefix: "h", file: file, line: line)
+    try assertCustomCompletion(
+      "five", shell: shell, prefix: "j", file: file, line: line)
 
     XCTAssertThrowsError(
       try assertCustomCompletion("--bad", shell: shell, file: file, line: line))