Skip to content

Support async custom completion closures #782

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 4, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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\
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 '(\
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ extension [ParsableCommand.Type] {
nil
)

case .custom:
case .custom, .customAsync:
return (
"{\(customCompleteFunctionName) \(arg.customCompletionCall(self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
nil
Expand Down
12 changes: 12 additions & 0 deletions Sources/ArgumentParser/Parsable Properties/CompletionKind.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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])
}

Expand Down Expand Up @@ -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])`,
Expand Down
120 changes: 93 additions & 27 deletions Sources/ArgumentParser/Parsing/CommandParser.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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 {
Expand Down
2 changes: 2 additions & 0 deletions Sources/ArgumentParser/Usage/DumpHelpGenerator.swift
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,8 @@ extension ArgumentInfoV0.CompletionKindV0 {
self = .shellCommand(command: command)
case .custom(_):
self = .custom
case .customAsync(_):
self = .customAsync
case .customDeprecated(_):
self = .customDeprecated
}
Expand Down
2 changes: 2 additions & 0 deletions Sources/ArgumentParserToolInfo/ToolInfo.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions Tests/ArgumentParserUnitTests/CompletionScriptTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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))
Expand Down
Loading