From 3a208cfd84a826eedc7eb6ec6c67546a2c28b8b0 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Tue, 22 Jul 2025 08:30:31 -0400
Subject: [PATCH] Autocomplete axolotl CLI (#2955)

* static autocomplete script for axolotl cli

* use list of commands that should autocomplete yaml files

* make sure to chmod the autocomplete script as executable

* shellcheck and fix autocompletion of directory/sub-dirs

* more shellcheck fixes
---
 .axolotl-complete.bash | 41 +++++++++++++++++++++++++++++++++++++++++
 docker/Dockerfile      |  6 +++++-
 2 files changed, 46 insertions(+), 1 deletion(-)
 create mode 100644 .axolotl-complete.bash

diff --git a/.axolotl-complete.bash b/.axolotl-complete.bash
new file mode 100644
index 000000000..9a51399e6
--- /dev/null
+++ b/.axolotl-complete.bash
@@ -0,0 +1,41 @@
+#!/bin/bash
+
+_axolotl_completions() {
+    local cur prev
+    COMPREPLY=()
+    cur="${COMP_WORDS[COMP_CWORD]}"
+    prev="${COMP_WORDS[COMP_CWORD-1]}"
+
+    # If we're completing the first argument (the command)
+    if [[ $COMP_CWORD -eq 1 ]]; then
+        mapfile -t COMPREPLY < <(compgen -W "delinearize-llama4 fetch lm-eval merge-sharded-fsdp-weights quantize vllm-serve evaluate inference merge-lora preprocess train" -- "$cur")
+        return 0
+    fi
+
+    # Commands that should complete with directories and YAML files
+    local -a yaml_commands=("merge-sharded-fsdp-weights" "quantize" "vllm-serve" "evaluate" "inference" "merge-lora" "preprocess" "train")
+
+    # Check if previous word is in our list
+    if [[ " ${yaml_commands[*]} " =~ (^|[[:space:]])$prev($|[[:space:]]) ]]; then
+        # Use filename completion which handles directories properly
+        compopt -o filenames
+        mapfile -t COMPREPLY < <(compgen -f -- "$cur")
+
+        # Filter to only include directories and YAML files
+        local -a filtered=()
+        for item in "${COMPREPLY[@]}"; do
+            if [[ -d "$item" ]] || [[ "$item" == *.yaml ]] || [[ "$item" == *.yml ]]; then
+                filtered+=("$item")
+            fi
+        done
+        COMPREPLY=("${filtered[@]}")
+
+        return 0
+    fi
+
+    # Default: no completion
+    return 0
+}
+
+# Remove the -o nospace option - let filenames handle it
+complete -F _axolotl_completions axolotl
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 9bc34c387..116361dcd 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -31,7 +31,11 @@ RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
     pip install pytest && \
     pip cache purge
 
-# fix so that git fetch/pull from remote works
+# fix so that git fetch/pull from remote works with shallow clone
 RUN git config remote.origin.fetch "+refs/heads/*:refs/remotes/origin/*" && \
     git config --get remote.origin.fetch && \
     git config --global credential.helper store
+
+COPY .axolotl-complete.bash /root/.axolotl-complete.bash
+RUN chmod +x /root/.axolotl-complete.bash && \
+    echo 'source /root/.axolotl-complete.bash' >> ~/.bashrc