Add completion script generation (apple#123)

Support for generating shell completion scripts for `ParsableCommand`
types, with customization points for `ExpressibleByArgument` types and
individual arguments and options. Zsh and Bash are supported in this
initial release.

Documentation/07 Completion Scripts.md

@@ -0,0 +1,108 @@
+# Completion Scripts
+
+Generate customized completion scripts for your shell of choice.
+
+## Generating and Installing Completion Scripts
+
+Command-line tools that you build with `ArgumentParser` include a built-in option for generating completion scripts, with support for Bash and Z shell. To generate completions, run your command with the `--generate-completion-script` flag to generate completions for the autodetected shell, or with a value to generate completions for a specific shell.
+
+```
+$ example --generate-completion-script bash
+#compdef example
+local context state state_descr line
+_example_commandname="example"
+typeset -A opt_args
+
+_example() {
+    integer ret=1
+    local -a args
+    ...
+}
+
+_example
+```
+
+The correct method of installing a completion script depends on your shell and your configuration.
+
+### Installing Zsh Completions
+
+If you have [`oh-my-zsh`](https://ohmyz.sh) installed, you already have a directory of automatically loading completion scripts — `.oh-my-zsh/completions`. Copy your new completion script to that directory.
+
+Without `oh-my-zsh`, you'll need to add a path for completion scripts to your function path, and turn on completion script autoloading. First, add these lines to `~/.zshrc`:
+
+```
+fpath=(~/.zsh/completion $fpath)
+autoload -U compinit
+compinit
+```
+
+Next, create a directory at `~/.zsh/completion` and copy the completion script to the new directory.
+
+### Installing Bash Completions
+
+If you have [`bash-completion`](https://github.com/scop/bash-completion) installed, you can just copy your new completion script to the `/usr/local/etc/bash_completion.d` directory.
+
+Without `bash-completion`, you'll need to source the completion script directly. Copy it to a directory such as `~/.bash_completions/`, and then add the following line to `~/.bash_profile` or `~/.bashrc`:
+
+```
+source ~/.bash_completions/example.bash
+```
+
+## Customizing Completions
+
+`ArgumentParser` provides default completions for any types that it can. For example, an `@Option` property that is a `CaseIterable` type will automatically have the correct values as completion suggestions.
+
+When declaring an option or argument, you can customize the completions that are offered by specifying a `CompletionKind`. With this completion kind you can specify that the value should be a file, a directory, or one of a list of strings:
+
+```swift
+struct Example: ParsableCommand {
+    @Option(help: "The file to read from.", completion: .file())
+    var input: String
+
+    @Option(help: "The output directory.", completion: .directory)
+    var outputDir: String
+
+    @Option(help: "The preferred file format.", completion: .list(["markdown", "rst"]))
+    var format: String
+
+    enum CompressionType: String, CaseIterable, ExpressibleByArgument {
+        case zip, gzip
+    }
+
+    @Option(help: "The compression type to use.")
+    var compression: CompressionType
+}
+```
+
+The generated completion script will suggest only file names for the `--input` option, only directory names for `--output-dir`, and only the strings `markdown` and `rst` for `--format`. The `--compression` option uses the default completions for a `CaseIterable` type, so the completion script will suggest `zip` and `gzip`.
+
+You can define the default completion kind for custom `ExpressibleByArgument` types by implementing `static var defaultCompletionKind: CompletionKind`. For example, any arguments or options with this `File` type will automatically use files for completions:
+
+```swift
+struct File: Hashable, ExpressibleByArgument {
+    var path: String
+
+    init?(argument: String) {
+        self.path = argument
+    }
+
+    static var defaultCompletionKind: CompletionKind {
+        .file()
+    }
+}
+```
+
+For even more control over the suggested completions, you can specify a function that will be called during completion by using the `.custom` completion kind.
+
+```swift
+func listExecutables(_ arguments: [String]) -> [String] {
+    // Generate the list of executables in the current directory
+}
+
+struct SwiftRun {
+    @Option(help: "The target to execute.", completion: .custom(listExecutables))
+    var target: String?
+}
+```
+
+In this example, when a user requests completions for the `--target` option, the completion script runs the `SwiftRun` command-line tool with a special syntax, calling the `listExecutables` function with an array of the arguments given so far.

Examples/math/main.swift

@@ -93,7 +93,7 @@ extension Math.Statistics {
             abstract: "Print the average of the values.",
             version: "1.5.0-alpha")
 
-        enum Kind: String, ExpressibleByArgument {
+        enum Kind: String, ExpressibleByArgument, CaseIterable {
             case mean, median, mode
         }
 
@@ -126,7 +126,7 @@ extension Math.Statistics {
             let sorted = values.sorted()
             let mid = sorted.count / 2
             if sorted.count.isMultiple(of: 2) {
-                return sorted[mid - 1] + sorted[mid] / 2
+                return (sorted[mid - 1] + sorted[mid]) / 2
             } else {
                 return sorted[mid]
             }
@@ -186,6 +186,12 @@ extension Math.Statistics {
         static var configuration = CommandConfiguration(
             abstract: "Print the quantiles of the values (TBD).")
 
+        @Argument(help: .hidden, completion: .list(["alphabet", "alligator", "branch", "braggart"]))
+        var oneOfFour: String?
+
+        @Argument(help: .hidden, completion: .custom { _ in ["alabaster", "breakfast", "crunch", "crash"] })
+        var customArg: String?
+
         @Argument(help: "A group of floating-point values to operate on.")
         var values: [Double] = []
 
@@ -199,6 +205,20 @@ extension Math.Statistics {
         @Option(help: .hidden)
         var testCustomExitCode: Int32?
 
+        // These args are for testing custom completion scripts:
+        @Option(help: .hidden, completion: .file(extensions: ["txt", "md"]))
+        var file: String?
+        @Option(help: .hidden, completion: .directory)
+        var directory: String?
+
+        @Option(
+          help: .hidden,
+          completion: .shellCommand("head -100 /usr/share/dict/words | tail -50"))
+        var shell: String?
+
+        @Option(help: .hidden, completion: .custom(customCompletion))
+        var custom: String?
+
         func validate() throws {
             if testSuccessExitCode {
                 throw ExitCode.success
@@ -219,4 +239,10 @@ extension Math.Statistics {
     }
 }
 
+func customCompletion(_ s: [String]) -> [String] {
+  return (s.last ?? "").starts(with: "a")
+    ? ["aardvark", "aaaaalbert"]
+    : ["hello", "helicopter", "heliotrope"]
+}
+
 Math.main()

Sources/ArgumentParser/CMakeLists.txt

@@ -1,6 +1,11 @@
 add_library(ArgumentParser
+  Completions/BashCompletionsGenerator.swift
+  Completions/CompletionsGenerator.swift
+  Completions/ZshCompletionsGenerator.swift
+
   "Parsable Properties/Argument.swift"
   "Parsable Properties/ArgumentHelp.swift"
+  "Parsable Properties/CompletionKind.swift"
   "Parsable Properties/Errors.swift"
   "Parsable Properties/Flag.swift"
   "Parsable Properties/NameSpecification.swift"

Sources/ArgumentParser/Completions/BashCompletionsGenerator.swift

@@ -0,0 +1,214 @@
+//===----------------------------------------------------------*- swift -*-===//
+//
+// This source file is part of the Swift Argument Parser open source project
+//
+// Copyright (c) 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+struct BashCompletionsGenerator {
+  /// Generates a Bash completion script for the given command.
+  static func generateCompletionScript(_ type: ParsableCommand.Type) -> String {
+    // TODO: Add a check to see if the command is installed where we expect?
+    let initialFunctionName = [type].completionFunctionName()
+    return """
+    #!/bin/bash
+
+    \(generateCompletionFunction([type]))
+
+    complete -F \(initialFunctionName) \(type._commandName)
+    """
+  }
+
+  /// 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()
+
+    // The root command gets a different treatment for the parsing index.
+    let isRootCommand = commands.count == 1
+    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
+    if !subcommands.isEmpty && isRootCommand {
+      subcommands.append(HelpCommand.self)
+    }
+
+    // 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 }
+      // FIXME: These shouldn't be hard-coded, since they're overridable
+      + ["-h", "--help"]
+
+    // Generate additional top-level completions — these are completion lists
+    // or custom function-based word lists from positional arguments.
+    let additionalCompletions = generateArgumentCompletions(commands)
+
+    // Start building the resulting function code.
+    var result = "\(functionName)() {\n"
+
+    // The function that represents the root command has some additional setup
+    // that other command functions don't need.
+    if isRootCommand {
+      result += """
+        cur="${COMP_WORDS[COMP_CWORD]}"
+        prev="${COMP_WORDS[COMP_CWORD-1]}"
+        COMPREPLY=()
+
+        """.indentingEachLine(by: 4)
+    }
+
+    // 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"
+    for line in additionalCompletions {
+      result += "    opts=\"$opts \(line)\"\n"
+    }
+
+    result += """
+        if [[ $COMP_CWORD == "\(dollarOne)" ]]; then
+            COMPREPLY=( $(compgen -W "$opts" -- "$cur") )
+            return
+        fi
+
+    """
+
+    // Generate the case pattern-matching statements for option values.
+    // If there aren't any, skip the case block altogether.
+    let optionHandlers = generateOptionHandlers(commands)
+    if !optionHandlers.isEmpty {
+      result += """
+      case $prev in
+      \(optionHandlers.indentingEachLine(by: 4))
+      esac
+      """.indentingEachLine(by: 4) + "\n"
+    }
+
+    // Build out completions for the subcommands.
+    if !subcommands.isEmpty {
+      // Subcommands have their own case statement that delegates out to
+      // the subcommand completion functions.
+      result += "    case ${COMP_WORDS[\(dollarOne)]} in\n"
+      for subcommand in subcommands {
+        result += """
+          (\(subcommand._commandName))
+              \(functionName)_\(subcommand._commandName) \(subcommandArgument)
+              return
+              ;;
+          
+          """
+          .indentingEachLine(by: 8)
+      }
+      result += "    esac\n"
+    }
+
+    // Finish off the function.
+    result += """
+        COMPREPLY=( $(compgen -W "$opts" -- "$cur") )
+    }
+
+    """
+
+    return result +
+      subcommands
+        .map { generateCompletionFunction(commands + [$0]) }
+        .joined()
+  }
+
+  /// Returns the option and flag names that can be top-level completions.
+  fileprivate static func generateArgumentWords(_ commands: [ParsableCommand.Type]) -> [String] {
+    ArgumentSet(commands.last!)
+      .flatMap { $0.bashCompletionWords() }
+  }
+
+  /// 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!)
+      .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:
+          // Generate a call back into the command to retrieve a completions list
+          let commandName = commands.first!._commandName
+          let subcommandNames = commands.dropFirst().map { $0._commandName }.joined(separator: " ")
+          // TODO: Make this work for @Arguments
+          let argumentName = arg.preferredNameForSynopsis?.synopsisString
+                ?? arg.help.keys.first?.rawValue ?? "---"
+
+          return """
+            $(\(commandName) ---completion \(subcommandNames) -- \(argumentName) "$COMP_WORDS")
+            """
+        }
+      }
+  }
+
+  /// Returns the case-matching statements for supplying completions after an option or flag.
+  fileprivate static func generateOptionHandlers(_ commands: [ParsableCommand.Type]) -> String {
+    ArgumentSet(commands.last!)
+      .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))
+            return
+        ;;
+        """
+      }
+      .joined(separator: "\n")
+  }
+}
+
+extension ArgumentDefinition {
+  /// Returns the different completion names for this argument.
+  fileprivate func bashCompletionWords() -> [String] {
+    names.map { $0.synopsisString }
+  }
+
+  /// Returns the bash completions that can follow this argument's `--name`.
+  fileprivate func bashValueCompletion(_ commands: [ParsableCommand.Type]) -> String {
+    switch completion.kind {
+    case .default:
+      return ""
+
+    case .file(_):
+      // TODO: Use '_filedir' when available
+      // FIXME: Use the extensions array
+      return #"COMPREPLY=( $(compgen -f -- "$cur") )"#
+
+    case .directory:
+      return #"COMPREPLY=( $(compgen -d -- "$cur") )"#
+
+    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
+      let commandName = commands.first!._commandName      
+      return #"COMPREPLY=( $(compgen -W "$(\#(commandName) \#(customCompletionCall(commands)) "$COMP_WORDS")" -- "$cur") )"#
+    }
+  }
+}