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

Issue:
Permit asynchronous code to be given to the custom completion type
<!--
    Thanks for contributing to the Swift Argument Parser!

    Before you submit your issue, please replace the paragraph
    below with information about your proposed feature.
-->

At the moment the `custom` case of the `CompletionKind` enumeration supports being given a custom closure to generate completions on the command line. However, since the provided closure must be synchronous, it is difficult to incorporate functionality that requires swift's concurrency features. The lack of this functionality prevents certain tasks from being possible with these completions and causes other tasks to be possible, but only in a non-thread safe way. Examples of this include working with databases (such as those managed by Apple's Core Data framework).

Example:

```swift
struct SomeCommand: AsyncParseableCommand {

@Argument(help: "<Helpful help>", completion: .custom({ arguments in
return try? await completionGenerator(arguments: arguments) ?? []
}))
var argument: String

mutating func run() async throws {
// Code that makes the command Run
}
}

@MainActor
private func completionGenerator(arguments: [String]) async throws -> [String] {
// Code that generates completions that has some asynchronous system in use and/or must be run on the main actor.
}
```

Pull request description:
Support async custom completion closures.

Resolve #555

### 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

Patch:
diff --git a/Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift b/Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift
index 87bf041ed..af6407e2d 100644
--- a/Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift
+++ b/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\
diff --git a/Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift b/Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift
index 6a1b56b59..525a14fd1 100644
--- a/Sources/ArgumentParser/Completions/FishCompletionsGenerator.swift
+++ b/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 '(\
diff --git a/Sources/ArgumentParser/Completions/ZshCompletionsGenerator.swift b/Sources/ArgumentParser/Completions/ZshCompletionsGenerator.swift
index cedc167b7..5e2001cad 100644
--- a/Sources/ArgumentParser/Completions/ZshCompletionsGenerator.swift
+++ b/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
diff --git a/Sources/ArgumentParser/Parsable Properties/CompletionKind.swift b/Sources/ArgumentParser/Parsable Properties/CompletionKind.swift
index e6ca38037..82739618b 100644
--- a/Sources/ArgumentParser/Parsable Properties/CompletionKind.swift
+++ b/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])`,
diff --git a/Sources/ArgumentParser/Parsing/CommandParser.swift b/Sources/ArgumentParser/Parsing/CommandParser.swift
index f45939c6f..97402bef7 100644
--- a/Sources/ArgumentParser/Parsing/CommandParser.swift
+++ b/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 {
diff --git a/Sources/ArgumentParser/Usage/DumpHelpGenerator.swift b/Sources/ArgumentParser/Usage/DumpHelpGenerator.swift
index 438ae411b..487d372cc 100644
--- a/Sources/ArgumentParser/Usage/DumpHelpGenerator.swift
+++ b/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
     }
diff --git a/Sources/ArgumentParserToolInfo/ToolInfo.swift b/Sources/ArgumentParserToolInfo/ToolInfo.swift
index 808106a87..3fe9014fc 100644
--- a/Sources/ArgumentParserToolInfo/ToolInfo.swift
+++ b/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


Test patch:
diff --git a/Tests/ArgumentParserUnitTests/CompletionScriptTests.swift b/Tests/ArgumentParserUnitTests/CompletionScriptTests.swift
index 06a46c179..f2ac309fe 100644
--- a/Tests/ArgumentParserUnitTests/CompletionScriptTests.swift
+++ b/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))