Skip to content

Commit 74b180f

Browse files
rgoldbergrgoldberg
committed
Refactor zsh completions to use ToolInfoV0.
Signed-off-by: Ross Goldberg <[email protected]>
1 parent 5218cf0 commit 74b180f

File tree

4 files changed

+87
-130
lines changed

4 files changed

+87
-130
lines changed

‎Sources/ArgumentParser/Completions/CompletionsGenerator.swift

+5-51
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ struct CompletionsGenerator {
140140
CompletionShell._requesting.withLock { $0 = shell }
141141
switch shell {
142142
case .zsh:
143-
return [command].zshCompletionScript
143+
return ToolInfoV0(commandStack: [command]).zshCompletionScript
144144
case .bash:
145145
return ToolInfoV0(commandStack: [command]).bashCompletionScript
146146
case .fish:
@@ -151,56 +151,6 @@ struct CompletionsGenerator {
151151
}
152152
}
153153

154-
extension ArgumentDefinition {
155-
/// Returns a string with the arguments for the callback to generate custom completions for
156-
/// this argument.
157-
func customCompletionCall(_ commands: [ParsableCommand.Type]) -> String {
158-
let subcommandNames =
159-
commands.dropFirst().map { "\($0._commandName) " }.joined()
160-
let argumentName =
161-
names.preferredName?.synopsisString
162-
?? self.help.keys.first?.fullPathString
163-
?? "---"
164-
return "---completion \(subcommandNames)-- \(argumentName)"
165-
}
166-
}
167-
168-
extension ParsableCommand {
169-
fileprivate static var compositeCommandName: [String] {
170-
if let superCommandName = configuration._superCommandName {
171-
return [superCommandName]
172-
+ _commandName.split(separator: " ").map(String.init)
173-
} else {
174-
return _commandName.split(separator: " ").map(String.init)
175-
}
176-
}
177-
}
178-
179-
extension [ParsableCommand.Type] {
180-
/// Include default 'help' subcommand in nonempty subcommand list if & only if
181-
/// no help subcommand already exists.
182-
mutating func addHelpSubcommandIfMissing() {
183-
if !isEmpty && !contains(where: { $0._commandName == "help" }) {
184-
append(HelpCommand.self)
185-
}
186-
}
187-
}
188-
189-
extension Sequence where Element == ParsableCommand.Type {
190-
func completionFunctionName() -> String {
191-
"_"
192-
+ self.flatMap { $0.compositeCommandName }
193-
.uniquingAdjacentElements()
194-
.joined(separator: "_")
195-
}
196-
197-
var shellVariableNamePrefix: String {
198-
flatMap { $0.compositeCommandName }
199-
.joined(separator: "_")
200-
.shellEscapeForVariableName()
201-
}
202-
}
203-
204154
extension String {
205155
func shellEscapeForSingleQuotedString(iterationCount: UInt64 = 1) -> Self {
206156
iterationCount == 0
@@ -234,6 +184,10 @@ extension CommandInfoV0 {
234184
var completionFunctionPrefix: String {
235185
"__\(initialCommand)"
236186
}
187+
188+
var shellVariableNamePrefix: String {
189+
commandContext.joined(separator: "_").shellEscapeForVariableName()
190+
}
237191
}
238192

239193
extension ArgumentInfoV0 {

‎Sources/ArgumentParser/Completions/ZshCompletionsGenerator.swift

+77-74
Original file line numberDiff line numberDiff line change
@@ -9,56 +9,64 @@
99
//
1010
//===----------------------------------------------------------------------===//
1111

12-
extension [ParsableCommand.Type] {
13-
/// Generates a Zsh completion script for the given command.
12+
#if swift(>=6.0)
13+
internal import ArgumentParserToolInfo
14+
#else
15+
import ArgumentParserToolInfo
16+
#endif
17+
18+
extension ToolInfoV0 {
1419
var zshCompletionScript: String {
20+
command.zshCompletionScript
21+
}
22+
}
23+
24+
extension CommandInfoV0 {
25+
fileprivate var zshCompletionScript: String {
1526
// swift-format-ignore: NeverForceUnwrap
1627
// Preconditions:
1728
// - first must be non-empty for a zsh completion script to be of use.
1829
// - first is guaranteed non-empty in the one place where this computed var is used.
19-
let commandName = first!._commandName
20-
return """
21-
#compdef \(commandName)
30+
"""
31+
#compdef \(commandName)
2232
23-
\(completeFunctionName)() {
24-
local -ar non_empty_completions=("${@:#(|:*)}")
25-
local -ar empty_completions=("${(M)@:#(|:*)}")
26-
_describe '' non_empty_completions -- empty_completions -P $'\\'\\''
27-
}
33+
\(completeFunctionName)() {
34+
local -ar non_empty_completions=("${@:#(|:*)}")
35+
local -ar empty_completions=("${(M)@:#(|:*)}")
36+
_describe '' non_empty_completions -- empty_completions -P $'\\'\\''
37+
}
2838
29-
\(customCompleteFunctionName)() {
30-
local -a completions
31-
completions=("${(@f)"$("${command_name}" "${@}" "${command_line[@]}")"}")
32-
if [[ "${#completions[@]}" -gt 1 ]]; then
33-
\(completeFunctionName) "${completions[@]:0:-1}"
34-
fi
35-
}
39+
\(customCompleteFunctionName)() {
40+
local -a completions
41+
completions=("${(@f)"$("${command_name}" "${@}" "${command_line[@]}")"}")
42+
if [[ "${#completions[@]}" -gt 1 ]]; then
43+
\(completeFunctionName) "${completions[@]:0:-1}"
44+
fi
45+
}
3646
37-
\(cursorIndexInCurrentWordFunctionName)() {
38-
if [[ -z "${QIPREFIX}${IPREFIX}${PREFIX}" ]]; then
39-
printf 0
40-
else
41-
printf %s "${#${(z)LBUFFER}[-1]}"
42-
fi
43-
}
47+
\(cursorIndexInCurrentWordFunctionName)() {
48+
if [[ -z "${QIPREFIX}${IPREFIX}${PREFIX}" ]]; then
49+
printf 0
50+
else
51+
printf %s "${#${(z)LBUFFER}[-1]}"
52+
fi
53+
}
4454
45-
\(completionFunctions)\
46-
\(completionFunctionName())
47-
"""
55+
\(completionFunctions)\
56+
\(completionFunctionName)
57+
"""
4858
}
4959

5060
private var completionFunctions: String {
51-
guard let type = last else { return "" }
52-
let functionName = completionFunctionName()
53-
let isRootCommand = count == 1
61+
let functionName = completionFunctionName
5462

55-
let argumentSpecsAndSetupScripts = argumentsForHelp(visibility: .default)
56-
.compactMap { argumentSpecAndSetupScript($0) }
63+
let argumentSpecsAndSetupScripts = (arguments ?? []).compactMap {
64+
argumentSpecAndSetupScript($0)
65+
}
5766
var argumentSpecs = argumentSpecsAndSetupScripts.map(\.argumentSpec)
5867
let setupScripts = argumentSpecsAndSetupScripts.compactMap(\.setupScript)
5968

60-
var subcommands = type.configuration.subcommands
61-
.filter { $0.configuration.shouldDisplay }
69+
let subcommands = (subcommands ?? []).filter(\.shouldDisplay)
6270

6371
let subcommandHandler: String
6472
if subcommands.isEmpty {
@@ -67,17 +75,13 @@ extension [ParsableCommand.Type] {
6775
argumentSpecs.append("'(-): :->command'")
6876
argumentSpecs.append("'(-)*:: :->arg'")
6977

70-
if isRootCommand {
71-
subcommands.addHelpSubcommandIfMissing()
72-
}
73-
7478
subcommandHandler = """
7579
case "${state}" in
7680
command)
7781
local -ar subcommands=(
7882
\(
7983
subcommands.map { """
80-
'\($0._commandName):\($0.configuration.abstract.zshEscapeForSingleQuotedExplanation())'
84+
'\($0.commandName):\($0.abstract?.zshEscapeForSingleQuotedExplanation() ?? "")'
8185
"""
8286
}
8387
.joined(separator: "\n")
@@ -87,7 +91,7 @@ extension [ParsableCommand.Type] {
8791
;;
8892
arg)
8993
case "${words[1]}" in
90-
\(subcommands.map { $0._commandName }.joined(separator: "|")))
94+
\(subcommands.map(\.commandName).joined(separator: "|")))
9195
"\(functionName)_${words[1]}"
9296
;;
9397
esac
@@ -99,7 +103,7 @@ extension [ParsableCommand.Type] {
99103

100104
return """
101105
\(functionName)() {
102-
\(isRootCommand
106+
\((superCommands ?? []).isEmpty
103107
? """
104108
emulate -RL zsh -G
105109
setopt extendedglob nullglob numericglobsort
@@ -131,47 +135,50 @@ extension [ParsableCommand.Type] {
131135
return "${ret}"
132136
}
133137
134-
\(subcommands.map { (self + [$0]).completionFunctions }.joined())
138+
\(subcommands.map(\.completionFunctions).joined())
135139
"""
136140
}
137141

138142
private func argumentSpecAndSetupScript(
139-
_ arg: ArgumentDefinition
143+
_ arg: ArgumentInfoV0
140144
) -> (argumentSpec: String, setupScript: String?)? {
141-
guard arg.help.visibility.base == .default else { return nil }
145+
guard arg.shouldDisplay else { return nil }
142146

143147
let line: String
144-
switch arg.names.count {
148+
let names = arg.names ?? []
149+
switch names.count {
145150
case 0:
146151
line = ""
147152
case 1:
153+
// swift-format-ignore: NeverForceUnwrap
154+
// Preconditions: names has exactly one element.
148155
line = """
149-
\(arg.isRepeatableOption ? "*" : "")\(arg.names[0].synopsisString)\(arg.zshCompletionAbstract)
156+
\(arg.isRepeatableOption ? "*" : "")\(names.first!.commonCompletionSynopsisString())\(arg.completionAbstract)
150157
"""
151158
default:
152-
let synopses = arg.names.map { $0.synopsisString }
159+
let synopses = names.map { $0.commonCompletionSynopsisString() }
153160
line = """
154161
\(arg.isRepeatableOption ? "*" : "(\(synopses.joined(separator: " ")))")'\
155162
{\(synopses.joined(separator: ","))}\
156-
'\(arg.zshCompletionAbstract)
163+
'\(arg.completionAbstract)
157164
"""
158165
}
159166

160-
switch arg.update {
161-
case .unary:
167+
switch arg.kind {
168+
case .option, .positional:
162169
let (argumentAction, setupScript) = argumentActionAndSetupScript(arg)
163-
return ("'\(line):\(arg.valueName):\(argumentAction)'", setupScript)
164-
case .nullary:
170+
return ("'\(line):\(arg.valueName ?? ""):\(argumentAction)'", setupScript)
171+
case .flag:
165172
return ("'\(line)'", nil)
166173
}
167174
}
168175

169176
/// Returns the zsh "action" for an argument completion string.
170177
private func argumentActionAndSetupScript(
171-
_ arg: ArgumentDefinition
178+
_ arg: ArgumentInfoV0
172179
) -> (argumentAction: String, setupScript: String?) {
173-
switch arg.completion.kind {
174-
case .default:
180+
switch arg.completionKind {
181+
case .none:
175182
return ("", nil)
176183

177184
case .file(let extensions):
@@ -201,41 +208,37 @@ extension [ParsableCommand.Type] {
201208

202209
case .custom:
203210
return (
204-
"{\(customCompleteFunctionName) \(arg.customCompletionCall(self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
211+
"{\(customCompleteFunctionName) \(arg.commonCustomCompletionCall(command: self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
205212
nil
206213
)
207214

208215
case .customDeprecated:
209216
return (
210-
"{\(customCompleteFunctionName) \(arg.customCompletionCall(self))}",
217+
"{\(customCompleteFunctionName) \(arg.commonCustomCompletionCall(command: self))}",
211218
nil
212219
)
213220
}
214221
}
215222

216-
private func variableName(_ arg: ArgumentDefinition) -> String {
217-
guard let argName = arg.names.preferredName else {
223+
private func variableName(_ arg: ArgumentInfoV0) -> String {
224+
guard let argName = arg.preferredName else {
218225
return
219-
"\(shellVariableNamePrefix)_\(arg.valueName.shellEscapeForVariableName())"
226+
"\(shellVariableNamePrefix)_\(arg.valueName?.shellEscapeForVariableName() ?? "")"
220227
}
221228
return
222-
"\(argName.case == .long ? "__" : "_")\(shellVariableNamePrefix)_\(argName.valueString.shellEscapeForVariableName())"
229+
"\(argName.kind == .long ? "__" : "_")\(shellVariableNamePrefix)_\(argName.name.shellEscapeForVariableName())"
223230
}
224231

225232
private var completeFunctionName: String {
226-
// swift-format-ignore: NeverForceUnwrap
227-
// Precondition: first is guaranteed to be non-empty
228-
"__\(first!._commandName)_complete"
233+
"\(completionFunctionPrefix)_complete"
229234
}
230235

231236
private var customCompleteFunctionName: String {
232-
// swift-format-ignore: NeverForceUnwrap
233-
// Precondition: first is guaranteed to be non-empty
234-
"__\(first!._commandName)_custom_complete"
237+
"\(completionFunctionPrefix)_custom_complete"
235238
}
236239

237240
private var cursorIndexInCurrentWordFunctionName: String {
238-
"__\(first?._commandName ?? "")_cursor_index_in_current_word"
241+
"\(completionFunctionPrefix)_cursor_index_in_current_word"
239242
}
240243
}
241244

@@ -250,13 +253,13 @@ extension String {
250253
}
251254
}
252255

253-
extension ArgumentDefinition {
256+
extension ArgumentInfoV0 {
254257
/// - returns: `true` if `self` is an option and can be tab-completed multiple times in one command line.
255258
/// For example, `ssh` allows the `-L` option to be given multiple times, to establish multiple port forwardings.
256259
fileprivate var isRepeatableOption: Bool {
257260
guard
258-
case .named(_) = kind,
259-
help.options.contains(.isRepeating)
261+
[.flag, .option].contains(kind),
262+
isRepeating
260263
else { return false }
261264

262265
switch parsingStrategy {
@@ -265,8 +268,8 @@ extension ArgumentDefinition {
265268
}
266269
}
267270

268-
fileprivate var zshCompletionAbstract: String {
269-
guard !help.abstract.isEmpty else { return "" }
270-
return "[\(help.abstract.zshEscapeForSingleQuotedExplanation())]"
271+
fileprivate var completionAbstract: String {
272+
guard let abstract, !abstract.isEmpty else { return "" }
273+
return "[\(abstract.zshEscapeForSingleQuotedExplanation())]"
271274
}
272275
}

‎Tests/ArgumentParserExampleTests/Snapshots/testMathZshCompletionScript().zsh

+2-2
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,8 @@ _math_stats_quantiles() {
156156
local -ar math_stats_quantiles_one_of_four=('alphabet' 'alligator' 'branch' 'braggart')
157157
local -ar arg_specs=(
158158
':one-of-four:{__math_complete "${math_stats_quantiles_one_of_four[@]}"}'
159-
':custom-arg:{__math_custom_complete ---completion stats quantiles -- customArg "${current_word_index}" "$(__math_cursor_index_in_current_word)"}'
160-
':custom-deprecated-arg:{__math_custom_complete ---completion stats quantiles -- customDeprecatedArg}'
159+
':custom-arg:{__math_custom_complete ---completion stats quantiles -- positional@1 "${current_word_index}" "$(__math_cursor_index_in_current_word)"}'
160+
':custom-deprecated-arg:{__math_custom_complete ---completion stats quantiles -- positional@2}'
161161
':values:'
162162
'--file:file:_files -g '\''*.txt *.md'\'''
163163
'--directory:directory:_files -/'

‎Tests/ArgumentParserUnitTests/Snapshots/testBase_Zsh().zsh

+3-3
Original file line numberDiff line numberDiff line change
@@ -56,8 +56,8 @@ _base-test() {
5656
'*--kind-counter'
5757
'*--rep1:rep1:'
5858
'*'{-r,--rep2}':rep2:'
59-
':argument:{__base-test_custom_complete ---completion -- argument "${current_word_index}" "$(__base-test_cursor_index_in_current_word)"}'
60-
':nested-argument:{__base-test_custom_complete ---completion -- nested.nestedArgument "${current_word_index}" "$(__base-test_cursor_index_in_current_word)"}'
59+
':argument:{__base-test_custom_complete ---completion -- positional@0 "${current_word_index}" "$(__base-test_cursor_index_in_current_word)"}'
60+
':nested-argument:{__base-test_custom_complete ---completion -- positional@1 "${current_word_index}" "$(__base-test_cursor_index_in_current_word)"}'
6161
'(-h --help)'{-h,--help}'[Show help information.]'
6262
'(-): :->command'
6363
'(-)*:: :->arg'
@@ -98,7 +98,7 @@ _base-test_escaped-command() {
9898
local -i ret=1
9999
local -ar arg_specs=(
100100
'--one[Escaped chars: '\''\[\]\\.]:one:'
101-
':two:{__base-test_custom_complete ---completion escaped-command -- two "${current_word_index}" "$(__base-test_cursor_index_in_current_word)"}'
101+
':two:{__base-test_custom_complete ---completion escaped-command -- positional@0 "${current_word_index}" "$(__base-test_cursor_index_in_current_word)"}'
102102
'(-h --help)'{-h,--help}'[Show help information.]'
103103
)
104104
_arguments -w -s -S : "${arg_specs[@]}" && ret=0

0 commit comments

Comments
 (0)