auto-completions for bb.cli/dispatch

src/babashka/cli.cljc

@@ -4,7 +4,8 @@
    #?(:clj [clojure.edn :as edn]
       :cljs [cljs.reader :as edn])
    [babashka.cli.internal :as internal]
-   [clojure.string :as str])
+   [clojure.string :as str]
+   [clojure.set :as set])
   #?(:clj (:import (clojure.lang ExceptionInfo))))
 
 #?(:clj (set! *warn-on-reflection* true))
@@ -589,8 +590,116 @@
           {} table))
 
 (comment
-  (table->tree [{:cmds [] :fn identity}])
- )
+  (table->tree [{:cmds [] :fn identity}]))
+
+;; completion
+(defn format-long-opt [k]
+  (str "--" (kw->str k)))
+(defn format-short-opt [k]
+  (str "-" (kw->str k)))
+
+(defn possibilities [cmd-tree]
+  (concat (keys (:cmd cmd-tree))
+          (map format-long-opt (keys (:spec cmd-tree)))
+          (map format-short-opt (keep :alias (vals (:spec cmd-tree))))))
+
+(defn true-prefix? [prefix s]
+  (and (< (count prefix) (count s))
+       (str/starts-with? s prefix)))
+
+(defn second-to-last [xs]
+  (when (>= (count xs) 2) (nth xs (- (count xs) 2))))
+
+(def possible-values (constantly []))
+
+(defn strip-prefix [prefix s]
+  (if (str/starts-with? s prefix)
+    (subs s (count prefix))
+    s))
+
+(defn bool-opt? [o spec]
+  (let [long-opt? (str/starts-with? o "--")
+        opt-kw (if long-opt?
+                 (keyword (strip-prefix "--" o))
+                 (some (fn [[k v]] (when (= (keyword (strip-prefix "-" o)) (:alias v)) k)) spec))]
+    (= :boolean (get-in spec [opt-kw :coerce]))))
+
+(defn is-gnu-option? [s]
+  (and s (str/starts-with? s "-")))
+
+(defn complete-tree
+  "given a CLI spec in tree form and input as a list of tokens,
+  returns possible tokens to complete the input"
+  [cmd-tree input]
+  (let [[head & tail] input
+        head (or head "")
+        subtree (get-in cmd-tree [:cmd head])]
+    (if (and subtree (first tail))
+      ;; matching command -> descend tree
+      (complete-tree subtree tail)
+      (if (is-gnu-option? head)
+        (let [{:keys [args opts err]} (try (parse-args input cmd-tree)
+                                           (catch clojure.lang.ExceptionInfo _ {:err :error}))]
+          (if (and args (not (str/blank? (first args))))
+            ;; parsed/consumed options and still have args left -> descend tree
+            (complete-tree cmd-tree args)
+            ;; no more args -> last input is (part of) an opt or opt value or empty string
+            (let [to-complete (last input)
+                  previous-token (second-to-last input)]
+              (if (and (is-gnu-option? previous-token) (not (bool-opt? previous-token (:spec cmd-tree))))
+                ;; complete value
+                (possible-values previous-token)
+                (let [possible-commands (keys (:cmd cmd-tree))
+                      ;; don't suggest options which we already have parsed
+                      possible-options (set/difference (set (keys (:spec cmd-tree))) (set (keys opts)))
+                      ;; generate string representation of possible options
+                      possible-completions (concat possible-commands
+                                                   (map format-long-opt possible-options)
+                                                   (keep (fn [option-name]
+                                                           (when-let [alias (get-in cmd-tree [:spec option-name :alias])]
+                                                             (format-short-opt alias)))
+                                                         possible-options))]
+                  (filter (partial true-prefix? to-complete) possible-completions))))))
+        (filter (partial true-prefix? head) (possibilities cmd-tree))))))
+
+(defn complete [cmd-table input]
+  (complete-tree (table->tree cmd-table) input))
+
+
+(defn generate-completion-shell-snippet [type program-name]
+  (case type
+    :bash (format "_babashka_cli_dynamic_completion()
+{
+    source <( \"$1\" --babashka.cli/complete \"bash\" \"${COMP_WORDS[*]// / }\" )
+}
+complete -o nosort -F _babashka_cli_dynamic_completion %s
+" program-name)
+    :zsh (format "#compdef %s
+source <( \"${words[1]}\" --babashka.cli/complete \"zsh\" \"${words[*]// / }\" )
+" program-name)
+    :fish (format "function _babashka_cli_dynamic_completion
+    set --local COMP_LINE (commandline --cut-at-cursor)
+    %s --babashka.cli/complete fish $COMP_LINE
+end
+complete --command %s --no-files --arguments \"(_babashka_cli_dynamic_completion)\"
+" program-name program-name)))
+
+(defn print-completion-shell-snippet [type program-name]
+  (print (generate-completion-shell-snippet type program-name)))
+
+(defn format-completion [shell {:keys [completion description]}]
+  (case shell
+    :bash (format "COMPREPLY+=( \"%s\" )" completion)
+    :zsh (str "compadd" (when description (str " -x \"" description "\"")) " -- " completion)
+    :fish completion))
+
+(defn print-completions [shell tree cmdline]
+  (let [[_program-name & to-complete] (str/split (str/triml cmdline) #" +" -1)
+        completions (complete-tree tree to-complete)]
+    (doseq [completion completions]
+      (println (format-completion shell {:completion completion})))))
+
+;; dispatch
 
 (defn- deep-merge [a b]
   (reduce (fn [acc k] (update acc k (fn [v]
@@ -653,21 +762,28 @@
   ([tree args]
    (dispatch-tree tree args nil))
   ([tree args opts]
-   (let [{:as res :keys [cmd-info error wrong-input available-commands]}
-         (dispatch-tree' tree args opts)
-         error-fn* (or (:error-fn opts)
-                       (fn [{:keys [msg] :as data}]
-                         (throw (ex-info msg data))))
-         error-fn (fn [data]
-                    (-> {;; :tree tree
-                         :type :org.babashka/cli
-                         :wrong-input wrong-input :all-commands available-commands}
-                        (merge data)
-                        error-fn*))]
-     (case error
-       (:no-match :input-exhausted)
-       (error-fn {:cause error :opts (:opts res)})
-       nil ((:fn cmd-info) (dissoc res :cmd-info))))))
+   (let [command-name (get-in opts [:completion :command])
+         [opt shell cmdline] args]
+     (case opt
+       "--babashka.cli/completion-snippet"
+       (print-completion-shell-snippet (keyword shell) command-name)
+       "--babashka.cli/complete"
+       (print-completions (keyword shell) tree cmdline)
+       (let [{:as res :keys [cmd-info error wrong-input available-commands]}
+             (dispatch-tree' tree args opts)
+             error-fn* (or (:error-fn opts)
+                           (fn [{:keys [msg] :as data}]
+                             (throw (ex-info msg data))))
+             error-fn (fn [data]
+                        (-> { ;; :tree tree
+                             :type :org.babashka/cli
+                             :wrong-input wrong-input :all-commands available-commands}
+                            (merge data)
+                            error-fn*))]
+         (case error
+           (:no-match :input-exhausted)
+           (error-fn {:cause error :opts (:opts res)})
+           nil ((:fn cmd-info) (dissoc res :cmd-info))))))))
 
 (defn dispatch
   "Subcommand dispatcher.

test/babashka/cli/completion_test.clj

@@ -0,0 +1,95 @@
+(ns babashka.cli.completion-test
+  (:require [babashka.cli :as cli :refer [complete]]
+            [clojure.java.io :as io]
+            [clojure.test :refer :all]))
+
+(def cmd-table
+  [{:cmds ["foo"] :spec {:foo-opt {:coerece :string
+                                   :alias :f}
+                         :foo-opt2 {:coerece :string}
+                         :foo-flag {:coerce :boolean
+                                    :alias :l}}}
+   {:cmds ["foo" "bar"] :spec {:bar-opt {:coerce :keyword}
+                               :bar-flag {:coerce :boolean}}}
+   {:cmds ["bar"]}
+   {:cmds ["bar-baz"]}])
+
+(deftest completion-test
+  (testing "complete commands"
+    (is (= #{"foo" "bar" "bar-baz"} (set (complete cmd-table [""]))))
+    (is (= #{"bar" "bar-baz"} (set (complete cmd-table ["ba"]))))
+    (is (= #{"bar-baz"} (set (complete cmd-table ["bar"]))))
+    (is (= #{"foo"} (set (complete cmd-table ["f"])))))
+
+  (testing "no completions for full command"
+    (is (= #{} (set (complete cmd-table ["foo"])))))
+
+  (testing "complete subcommands and options"
+    (is (= #{"bar" "-f" "--foo-opt" "--foo-opt2" "-l" "--foo-flag"} (set (complete cmd-table ["foo" ""])))))
+
+  (testing "complete suboption"
+    (is (= #{"-f" "--foo-opt" "--foo-opt2" "-l" "--foo-flag"} (set (complete cmd-table ["foo" "-"])))))
+
+  (testing "complete short-opt"
+    (is (= #{} (set (complete cmd-table ["foo" "-f"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" ""]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" "foo-val"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" "bar"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" "foo-flag"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" "foo-opt2"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" "123"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" ":foo"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "-f" "true"]))))
+    (is (= #{"bar" "--foo-opt2" "-l" "--foo-flag"} (set (complete cmd-table ["foo" "-f" "foo-val" ""])))))
+
+  (testing "complete option with same prefix"
+    (is (= #{"--foo-opt" "--foo-opt2" "--foo-flag"} (set (complete cmd-table ["foo" "--foo"]))))
+    (is (= #{"--foo-opt2"} (set (complete cmd-table ["foo" "--foo-opt"])))))
+
+  (testing "complete long-opt"
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt2"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" ""]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" "foo-val"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" "bar"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" "foo-flag"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" "foo-opt2"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" "123"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" ":foo"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-opt" "true"]))))
+    (is (= #{"bar" "--foo-opt2" "-l" "--foo-flag"} (set (complete cmd-table ["foo" "--foo-opt" "foo-val" ""])))))
+
+  (is (= #{"--foo-flag"} (set (complete cmd-table ["foo" "--foo-f"]))))
+
+  (testing "complete short flag"
+    (is (= #{} (set (complete cmd-table ["foo" "-l"]))))
+    (is (= #{"bar" "-f" "--foo-opt" "--foo-opt2"} (set (complete cmd-table ["foo" "-l" ""])))))
+
+  (testing "complete long flag"
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-flag"]))))
+    (is (= #{"bar" "-f" "--foo-opt" "--foo-opt2"} (set (complete cmd-table ["foo" "--foo-flag" ""])))))
+
+  (is (= #{"-f" "--foo-opt" "--foo-opt2"} (set (complete cmd-table ["foo" "--foo-flag" "-"]))))
+  (is (= #{"bar"} (set (complete cmd-table ["foo" "--foo-flag" "b"]))))
+
+  (testing "complete subcommand"
+    (is (= #{"--bar-opt" "--bar-flag"} (set (complete cmd-table ["foo" "--foo-flag" "bar" ""]))))
+    (is (= #{"--bar-opt" "--bar-flag"} (set (complete cmd-table ["foo" "--foo-flag" "bar" "-"]))))
+    (is (= #{"--bar-opt" "--bar-flag"} (set (complete cmd-table ["foo" "--foo-flag" "bar" "--"]))))
+    (is (= #{"--bar-opt" "--bar-flag"} (set (complete cmd-table ["foo" "--foo-flag" "bar" "--bar-"]))))
+    (is (= #{"--bar-opt"} (set (complete cmd-table ["foo" "--foo-flag" "bar" "--bar-o"]))))
+    (is (= #{} (set (complete cmd-table ["foo" "--foo-flag" "bar" "--bar-opt" "a"]))))
+    (is (= #{"--bar-flag"} (set (complete cmd-table ["foo" "--foo-flag" "bar" "--bar-opt" "bar-val" ""]))))))
+
+
+(deftest parse-opts-completion-test
+  (cli/parse-opts ["--babashka.cli/completion-snippet" "zsh"] {:complete true})
+  (cli/parse-opts ["--babashka.cli/complete" "zsh" "foo"] {:complete true}))
+
+(deftest dispatch-completion-test
+  (is (= (slurp (io/resource "resources/completion/completion.zsh")) (with-out-str (cli/dispatch cmd-table ["--babashka.cli/completion-snippet" "zsh"] {:completion {:command "myprogram"}}))))
+  (is (= (slurp (io/resource "resources/completion/completion.bash")) (with-out-str (cli/dispatch cmd-table ["--babashka.cli/completion-snippet" "bash"] {:completion {:command "myprogram"}}))))
+  (is (= (slurp (io/resource "resources/completion/completion.fish")) (with-out-str (cli/dispatch cmd-table ["--babashka.cli/completion-snippet" "fish"] {:completion {:command "myprogram"}}))))
+
+  (is (= "compadd -- foo\n" (with-out-str (cli/dispatch cmd-table ["--babashka.cli/complete" "zsh" "myprogram f"] {:completion {:command "myprogram"}}))))
+  (is (= "COMPREPLY+=( \"foo\" )\n" (with-out-str (cli/dispatch cmd-table ["--babashka.cli/complete" "bash" "myprogram f "] {:completion {:command "myprogram"}}))))
+  (is (= "foo\n" (with-out-str (cli/dispatch cmd-table ["--babashka.cli/complete" "fish" "myprogram f "] {:completion {:command "myprogram"}})))))

test/resources/completion/completion.bash

@@ -0,0 +1,5 @@
+_babashka_cli_dynamic_completion()
+{
+    source <( "$1" --babashka.cli/complete "bash" "${COMP_WORDS[*]// / }" )
+}
+complete -o nosort -F _babashka_cli_dynamic_completion myprogram

test/resources/completion/completion.fish

@@ -0,0 +1,5 @@
+function _babashka_cli_dynamic_completion
+    set --local COMP_LINE (commandline --cut-at-cursor)
+    myprogram --babashka.cli/complete fish $COMP_LINE
+end
+complete --command myprogram --no-files --arguments "(_babashka_cli_dynamic_completion)"

test/resources/completion/completion.zsh

@@ -0,0 +1,2 @@
+#compdef myprogram
+source <( "${words[1]}" --babashka.cli/complete "zsh" "${words[*]// / }" )