Skip to content

Commit e1cb25e

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

File tree

4 files changed

+99
-144
lines changed

4 files changed

+99
-144
lines changed

‎Sources/ArgumentParser/Completions/CompletionsGenerator.swift‎

Lines changed: 5 additions & 51 deletions
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‎

Lines changed: 79 additions & 78 deletions
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 -V '' non_empty_completions -- empty_completions -P $'\\'\\''
27-
}
33+
\(completeFunctionName)() {
34+
local -ar non_empty_completions=("${@:#(|:*)}")
35+
local -ar empty_completions=("${(M)@:#(|:*)}")
36+
_describe -V '' 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.zshEscapeForSingleQuotedDescribeCompletion()):\($0.configuration.abstract.shellEscapeForSingleQuotedString())'
84+
'\($0.commandName.zshEscapeForSingleQuotedDescribeCompletion()):\($0.abstract?.shellEscapeForSingleQuotedString() ?? "")'
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,52 +135,55 @@ 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:
146-
line = arg.help.options.contains(.isRepeating) ? "*" : ""
151+
line = arg.isRepeating ? "*" : ""
147152
case 1:
153+
// swift-format-ignore: NeverForceUnwrap
154+
// Preconditions: names has exactly one element.
148155
line = """
149-
\(arg.isRepeatingOption ? "*" : "")\(arg.names[0].synopsisString.zshEscapeForSingleQuotedOptionSpec())\(arg.zshCompletionAbstract)
156+
\(arg.isRepeatingOption ? "*" : "")\(names.first!.commonCompletionSynopsisString().zshEscapeForSingleQuotedOptionSpec())\(arg.completionAbstract)
150157
"""
151158
default:
152-
let synopses = arg.names.map {
153-
$0.synopsisString.zshEscapeForSingleQuotedOptionSpec()
159+
let synopses = names.map {
160+
$0.commonCompletionSynopsisString().zshEscapeForSingleQuotedOptionSpec()
154161
}
155162
line = """
156163
\(arg.isRepeatingOption ? "*" : "(\(synopses.joined(separator: " ")))")'\
157164
{\(synopses.joined(separator: ","))}\
158-
'\(arg.zshCompletionAbstract)
165+
'\(arg.completionAbstract)
159166
"""
160167
}
161168

162-
switch arg.update {
163-
case .unary:
169+
switch arg.kind {
170+
case .option, .positional:
164171
let (argumentAction, setupScript) = argumentActionAndSetupScript(arg)
165172
return (
166-
"'\(line):\(arg.valueName.zshEscapeForSingleQuotedOptionSpec()):\(argumentAction)'",
173+
"'\(line):\(arg.valueName?.zshEscapeForSingleQuotedOptionSpec() ?? ""):\(argumentAction)'",
167174
setupScript
168175
)
169-
case .nullary:
176+
case .flag:
170177
return ("'\(line)'", nil)
171178
}
172179
}
173180

174181
/// Returns the zsh "action" for an argument completion string.
175182
private func argumentActionAndSetupScript(
176-
_ arg: ArgumentDefinition
183+
_ arg: ArgumentInfoV0
177184
) -> (argumentAction: String, setupScript: String?) {
178-
switch arg.completion.kind {
179-
case .default:
185+
switch arg.completionKind {
186+
case .none:
180187
return ("", nil)
181188

182189
case .file(let extensions):
@@ -206,51 +213,46 @@ extension [ParsableCommand.Type] {
206213

207214
case .custom, .customAsync:
208215
return (
209-
"{\(customCompleteFunctionName) \(arg.customCompletionCall(self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
216+
"{\(customCompleteFunctionName) \(arg.commonCustomCompletionCall(command: self)) \"${current_word_index}\" \"$(\(cursorIndexInCurrentWordFunctionName))\"}",
210217
nil
211218
)
212219

213220
case .customDeprecated:
214221
return (
215-
"{\(customCompleteFunctionName) \(arg.customCompletionCall(self))}",
222+
"{\(customCompleteFunctionName) \(arg.commonCustomCompletionCall(command: self))}",
216223
nil
217224
)
218225
}
219226
}
220227

221-
private func variableName(_ arg: ArgumentDefinition) -> String {
222-
guard let argName = arg.names.preferredName else {
223-
return
224-
"\(shellVariableNamePrefix)_\(arg.valueName.shellEscapeForVariableName())"
228+
private func variableName(_ arg: ArgumentInfoV0) -> String {
229+
guard let argName = arg.preferredName else {
230+
return "_\(arg.valueName?.shellEscapeForVariableName() ?? "")"
225231
}
226232
return
227-
"\(argName.case == .long ? "__" : "_")\(shellVariableNamePrefix)_\(argName.valueString.shellEscapeForVariableName())"
233+
"\(argName.kind == .long ? "___" : "__")\(argName.name.shellEscapeForVariableName())"
228234
}
229235

230236
private var completeFunctionName: String {
231-
// swift-format-ignore: NeverForceUnwrap
232-
// Precondition: first is guaranteed to be non-empty
233-
"__\(first!._commandName)_complete"
237+
"\(completionFunctionPrefix)_complete"
234238
}
235239

236240
private var customCompleteFunctionName: String {
237-
// swift-format-ignore: NeverForceUnwrap
238-
// Precondition: first is guaranteed to be non-empty
239-
"__\(first!._commandName)_custom_complete"
241+
"\(completionFunctionPrefix)_custom_complete"
240242
}
241243

242244
private var cursorIndexInCurrentWordFunctionName: String {
243-
"__\(first?._commandName ?? "")_cursor_index_in_current_word"
245+
"\(completionFunctionPrefix)_cursor_index_in_current_word"
244246
}
245247
}
246248

247-
extension ArgumentDefinition {
249+
extension ArgumentInfoV0 {
248250
/// - returns: `true` if `self` is a flag or an option and can be tab-completed multiple times in one command line.
249251
/// For example, `ssh` allows the `-L` option to be given multiple times, to establish multiple port forwardings.
250252
fileprivate var isRepeatingOption: Bool {
251253
guard
252-
case .named(_) = kind,
253-
help.options.contains(.isRepeating)
254+
[.flag, .option].contains(kind),
255+
isRepeating
254256
else { return false }
255257

256258
switch parsingStrategy {
@@ -259,10 +261,9 @@ extension ArgumentDefinition {
259261
}
260262
}
261263

262-
fileprivate var zshCompletionAbstract: String {
263-
help.abstract.isEmpty
264-
? ""
265-
: "[\(help.abstract.zshEscapeForSingleQuotedOptionSpec())]"
264+
fileprivate var completionAbstract: String {
265+
guard let abstract, !abstract.isEmpty else { return "" }
266+
return "[\(abstract.zshEscapeForSingleQuotedOptionSpec())]"
266267
}
267268
}
268269

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -127,9 +127,9 @@ _math_stats() {
127127

128128
_math_stats_average() {
129129
local -i ret=1
130-
local -ar __math_stats_average_kind=('mean' 'median' 'mode')
130+
local -ar ___kind=('mean' 'median' 'mode')
131131
local -ar arg_specs=(
132-
'--kind[The kind of average to provide.]:kind:{__math_complete "${__math_stats_average_kind[@]}"}'
132+
'--kind[The kind of average to provide.]:kind:{__math_complete "${___kind[@]}"}'
133133
'*:values:'
134134
'--version[Show the version.]'
135135
'(-h --help)'{-h,--help}'[Show help information.]'
@@ -153,11 +153,11 @@ _math_stats_stdev() {
153153

154154
_math_stats_quantiles() {
155155
local -i ret=1
156-
local -ar math_stats_quantiles_one_of_four=('alphabet' 'alligator' 'branch' 'braggart')
156+
local -ar _one_of_four=('alphabet' 'alligator' 'branch' 'braggart')
157157
local -ar arg_specs=(
158-
':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}'
158+
':one-of-four:{__math_complete "${_one_of_four[@]}"}'
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 -/'

0 commit comments

Comments
 (0)