Skip to content

Commit

Permalink
Refactor Bash completions to use ToolInfo
Browse files Browse the repository at this point in the history
Rebases the implementation of the BashCompletionGenerator to use
ToolInfo from ArgumentParserToolInfo instead of digging through the
command structure. This helps us decouple the implementation of Argument
parsing from the generation of supplemental content such as docs,
man-pages, completion scripts, help menus and more.
  • Loading branch information
rauhul committed Feb 6, 2025
1 parent d3630e3 commit d66a07b
Show file tree
Hide file tree
Showing 9 changed files with 197 additions and 121 deletions.
1 change: 1 addition & 0 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ var package = Package(
.testTarget(
name: "ArgumentParserExampleTests",
dependencies: ["ArgumentParserTestHelpers"],
exclude: ["Snapshots"],
resources: [.copy("CountLinesTest.txt")]),
.testTarget(
name: "ArgumentParserGenerateDoccReferenceTests",
Expand Down
1 change: 1 addition & 0 deletions Package@swift-5.8.swift
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ var package = Package(
.testTarget(
name: "ArgumentParserExampleTests",
dependencies: ["ArgumentParserTestHelpers"],
exclude: ["Snapshots"],
resources: [.copy("CountLinesTest.txt")]),
.testTarget(
name: "ArgumentParserGenerateDoccReferenceTests",
Expand Down
199 changes: 114 additions & 85 deletions Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift
Original file line number Diff line number Diff line change
Expand Up @@ -9,47 +9,59 @@
//
//===----------------------------------------------------------------------===//

import ArgumentParserToolInfo

struct BashCompletionsGenerator {
/// Generates a Bash completion script for the given command.
static func generateCompletionScript(_ type: ParsableCommand.Type) -> String {
return ToolInfoV0(commandStack: [type]).bashCompletionScript()
}
}

extension ToolInfoV0 {
fileprivate func bashCompletionScript() -> String {
// TODO: Add a check to see if the command is installed where we expect?
let initialFunctionName = [type].completionFunctionName().makeSafeFunctionName
return """
#!/bin/bash
#!/bin/bash
\(self.command.bashCompletionFunction())
complete -F \(self.command.bashCompletionFunctionName()) \(self.command.commandName)
"""
}
}

\(generateCompletionFunction([type]))
extension CommandInfoV0 {
fileprivate func bashCommandContext() -> [String] {
return (self.superCommands ?? []) + [self.commandName]
}

complete -F \(initialFunctionName) \(type._commandName)
"""
fileprivate func bashCompletionFunctionName() -> String {
return "_" + self.bashCommandContext().joined(separator: "_").makeSafeFunctionName
}

/// Generates a Bash completion function for the last command in the given list.
fileprivate static func generateCompletionFunction(_ commands: [ParsableCommand.Type]) -> String {
let type = commands.last!
let functionName = commands.completionFunctionName().makeSafeFunctionName

/// Generates a Bash completion function.
fileprivate func bashCompletionFunction() -> String {
let functionName = self.bashCompletionFunctionName()

// The root command gets a different treatment for the parsing index.
let isRootCommand = commands.count == 1
let isRootCommand = (self.superCommands ?? []).count == 0
let dollarOne = isRootCommand ? "1" : "$1"
let subcommandArgument = isRootCommand ? "2" : "$(($1+1))"

// Include 'help' in the list of subcommands for the root command.
var subcommands = type.configuration.subcommands
.filter { $0.configuration.shouldDisplay }
if !subcommands.isEmpty && isRootCommand {
subcommands.append(HelpCommand.self)
}
let subcommands = (self.subcommands ?? [])
.filter { $0.shouldDisplay }

// Generate the words that are available at the "top level" of this
// command — these are the dash-prefixed names of options and flags as well
// as all the subcommand names.
let completionWords = generateArgumentWords(commands)
+ subcommands.map { $0._commandName }

let completionKeys = self.bashCompletionKeys() + subcommands.map { $0.commandName }

// Generate additional top-level completions — these are completion lists
// or custom function-based word lists from positional arguments.
let additionalCompletions = generateArgumentCompletions(commands)
let additionalCompletions = self.bashPositionalCompletions()

// Start building the resulting function code.
var result = "\(functionName)() {\n"

Expand All @@ -69,7 +81,7 @@ struct BashCompletionsGenerator {

// Start by declaring a local var for the top-level completions.
// Return immediately if the completion matching hasn't moved further.
result += " opts=\"\(completionWords.joined(separator: " "))\"\n"
result += " opts=\"\(completionKeys.joined(separator: " "))\"\n"
for line in additionalCompletions {
result += " opts=\"$opts \(line)\"\n"
}
Expand All @@ -84,7 +96,7 @@ struct BashCompletionsGenerator {

// Generate the case pattern-matching statements for option values.
// If there aren't any, skip the case block altogether.
let optionHandlers = generateOptionHandlers(commands)
let optionHandlers = self.bashOptionCompletions().joined(separator: "\n")
if !optionHandlers.isEmpty {
result += """
case $prev in
Expand All @@ -100,8 +112,8 @@ struct BashCompletionsGenerator {
result += " case ${COMP_WORDS[\(dollarOne)]} in\n"
for subcommand in subcommands {
result += """
(\(subcommand._commandName))
\(functionName)_\(subcommand._commandName) \(subcommandArgument)
(\(subcommand.commandName))
\(functionName)_\(subcommand.commandName) \(subcommandArgument)
return
;;
Expand All @@ -120,77 +132,100 @@ struct BashCompletionsGenerator {

return result +
subcommands
.map { generateCompletionFunction(commands + [$0]) }
.joined()
.map { $0.bashCompletionFunction() }
.joined()
}

/// Returns the option and flag names that can be top-level completions.
fileprivate static func generateArgumentWords(_ commands: [ParsableCommand.Type]) -> [String] {
commands
.argumentsForHelp(visibility: .default)
.flatMap { $0.bashCompletionWords() }
fileprivate func bashCompletionKeys() -> [String] {
var result = [String]()
for argument in self.arguments ?? [] {
// Skip hidden arguments.
guard argument.shouldDisplay else { continue }
result.append(contentsOf: argument.bashCompletionKeys())
}
return result
}

/// Returns additional top-level completions from positional arguments.
///
/// These consist of completions that are defined as `.list` or `.custom`.
fileprivate static func generateArgumentCompletions(_ commands: [ParsableCommand.Type]) -> [String] {
ArgumentSet(commands.last!, visibility: .default, parent: nil)
.compactMap { arg -> String? in
guard arg.isPositional else { return nil }

switch arg.completion.kind {
case .default, .file, .directory:
return nil
case .list(let list):
return list.joined(separator: " ")
case .shellCommand(let command):
return "$(\(command))"
case .custom:
return """
$("${COMP_WORDS[0]}" \(arg.customCompletionCall(commands)) "${COMP_WORDS[@]}")
"""
}
}
fileprivate func bashPositionalCompletions() -> [String] {
var result = [String]()
for argument in self.arguments ?? [] {
// Skip hidden arguments.
guard argument.shouldDisplay else { continue }
// Only select positional arguments.
guard argument.kind == .positional else { continue }
// Skip if no completions.
guard let completionValues = argument.bashPositionalCompletionValues(command: self) else { continue }
result.append(completionValues)
}
return result
}

/// Returns the case-matching statements for supplying completions after an option or flag.
fileprivate static func generateOptionHandlers(_ commands: [ParsableCommand.Type]) -> String {
ArgumentSet(commands.last!, visibility: .default, parent: nil)
.compactMap { arg -> String? in
let words = arg.bashCompletionWords()
if words.isEmpty { return nil }

// Flags don't take a value, so we don't provide follow-on completions.
if arg.isNullary { return nil }

return """
\(arg.bashCompletionWords().joined(separator: "|")))
\(arg.bashValueCompletion(commands).indentingEachLine(by: 4))
fileprivate func bashOptionCompletions() -> [String] {
var result = [String]()
for argument in self.arguments ?? [] {
// Skip hidden arguments.
guard argument.shouldDisplay else { continue }
// Flags don't take a value, so we don't provide follow-on completions.
guard argument.kind != .flag else { continue }
// Skip if no keys.
let keys = argument.bashCompletionKeys()
guard !keys.isEmpty else { continue }
// Skip if no completions.
guard let completionValues = argument.bashOptionCompletionValues(command: self) else { continue }
result.append("""
\(keys.joined(separator: "|")))
\(completionValues.indentingEachLine(by: 4))
return
;;
"""
}
.joined(separator: "\n")
""")
}
return result
}
}

extension ArgumentDefinition {
extension ArgumentInfoV0 {
/// Returns the different completion names for this argument.
fileprivate func bashCompletionWords() -> [String] {
return help.visibility.base == .default
? names.map { $0.synopsisString }
: []
fileprivate func bashCompletionKeys() -> [String] {
return (self.names ?? []).map { $0.commonCompletionSynopsisString() }
}

// FIXME: determine if this can be combined with bashOptionCompletionValues
fileprivate func bashPositionalCompletionValues(
command: CommandInfoV0
) -> String? {
precondition(self.kind == .positional)

switch self.completionKind {
case .none, .file, .directory:
// FIXME: this doesn't work
return nil
case .list(let list):
return list.joined(separator: " ")
case .shellCommand(let command):
return "$(\(command))"
case .custom:
// Generate a call back into the command to retrieve a completions list
return #"$("${COMP_WORDS[0]}" \#(self.commonCustomCompletionCall(command: command)) "${COMP_WORDS[@]}")"#
}
}

/// Returns the bash completions that can follow this argument's `--name`.
///
/// Uses bash-completion for file and directory values if available.
fileprivate func bashValueCompletion(_ commands: [ParsableCommand.Type]) -> String {
switch completion.kind {
case .default:
return ""

fileprivate func bashOptionCompletionValues(
command: CommandInfoV0
) -> String? {
precondition(self.kind == .option)

switch self.completionKind {
case .none:
return nil

case .file(let extensions) where extensions.isEmpty:
return """
if declare -F _filedir >/dev/null; then
Expand All @@ -203,7 +238,7 @@ extension ArgumentDefinition {
case .file(let extensions):
var safeExts = extensions.map { String($0.flatMap { $0 == "'" ? ["\\", "'"] : [$0] }) }
safeExts.append(contentsOf: safeExts.map { $0.uppercased() })

return """
if declare -F _filedir >/dev/null; then
\(safeExts.map { "_filedir '\($0)'" }.joined(separator:"\n "))
Expand All @@ -224,22 +259,16 @@ extension ArgumentDefinition {
COMPREPLY=( $(compgen -d -- "$cur") )
fi
"""

case .list(let list):
return #"COMPREPLY=( $(compgen -W "\#(list.joined(separator: " "))" -- "$cur") )"#

case .shellCommand(let command):
return "COMPREPLY=( $(\(command)) )"

case .custom:
// Generate a call back into the command to retrieve a completions list
return #"COMPREPLY=( $(compgen -W "$("${COMP_WORDS[0]}" \#(customCompletionCall(commands)) "${COMP_WORDS[@]}")" -- "$cur") )"#
return #"COMPREPLY=( $(compgen -W "$("${COMP_WORDS[0]}" \#(self.commonCustomCompletionCall(command: command)) "${COMP_WORDS[@]}")" -- "$cur") )"#
}
}
}

extension String {
var makeSafeFunctionName: String {
self.replacingOccurrences(of: "-", with: "_")
}
}
40 changes: 40 additions & 0 deletions Sources/ArgumentParser/Completions/CompletionsGenerator.swift
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
//
//===----------------------------------------------------------------------===//

import ArgumentParserToolInfo

/// A shell for which the parser can generate a completion script.
public struct CompletionShell: RawRepresentable, Hashable, CaseIterable {
public var rawValue: String
Expand Down Expand Up @@ -139,3 +141,41 @@ extension Sequence where Element == ParsableCommand.Type {
.joined(separator: "_")
}
}

extension String {
var makeSafeFunctionName: String {
self.replacingOccurrences(of: "-", with: "_")
}
}

extension ArgumentInfoV0 {
/// Returns a string with the arguments for the callback to generate custom
/// completions for this argument.
func commonCustomCompletionCall(command: CommandInfoV0) -> String {
let commandContext = (command.superCommands ?? []) + [command.commandName]
let subcommandNames = commandContext.dropFirst().joined(separator: " ")

let argumentName: String
switch self.kind {
case .positional:
let index = (command.arguments ?? [])
.filter { $0.kind == .positional }
.firstIndex(of: self)!
argumentName = "positional@\(index)"
default:
argumentName = self.preferredName!.commonCompletionSynopsisString()
}
return "---completion \(subcommandNames) -- \(argumentName)"
}
}

extension ArgumentInfoV0.NameInfoV0 {
func commonCompletionSynopsisString() -> String {
switch self.kind {
case .long:
return "--\(self.name)"
case .short, .longWithSingleDash:
return "-\(self.name)"
}
}
}
12 changes: 10 additions & 2 deletions Sources/ArgumentParser/Parsing/ArgumentSet.swift
Original file line number Diff line number Diff line change
Expand Up @@ -191,9 +191,17 @@ extension ArgumentSet {
}

func firstPositional(
withKey key: InputKey
withKey key: InputKey,
) -> ArgumentDefinition? {
first(where: { $0.help.keys.contains(key) })
return first(where: { $0.help.keys.contains(key) })
}

func positional(
at index: Int
) -> ArgumentDefinition? {
let arguments = self.content.filter { $0.isPositional }
guard arguments.count > index else { return nil }
return arguments[index]
}
}

Expand Down
Loading

0 comments on commit d66a07b

Please sign in to comment.