Fix Breeze2 autocomplete (#22695)

Breeze2 autocomplete did not work because we were using some old
way of adding it (via click-complete). Since then click has
native (and very well working) autocomplete support without any
external dependencies needed. It cannnot automatically generate
the completions but it is not needed either, because we can
store generated completion scripts in our repo.

We also move some imports to local and catch rich_click import
error to minimize dependencies needed to get autocomplete
working. Setup-autocomplete install (and upgrade if needed
click in case it needs to be used by autocomplete script.

Fixes: #21164
GitOrigin-RevId: 4fb929a60966e9cecbbb435efd375adcc4fff9d7

.pre-commit-config.yaml

@@ -94,7 +94,7 @@ repos:
           - --fuzzy-match-generates-todo
       - id: insert-license
         name: Add license for all shell files
-        exclude: ^\.github/.*$|^airflow/_vendor/
+        exclude: ^\.github/.*$|^airflow/_vendor/|^dev/breeze/autocomplete/.*$
         files: ^breeze$|^breeze-complete$|\.bash$|\.sh$
         args:
           - --comment-style
@@ -524,6 +524,7 @@ repos:
         language: docker_image
         entry: koalaman/shellcheck:v0.7.2 -x -a
         files: ^breeze$|^breeze-complete$|\.sh$|^hooks/build$|^hooks/push$|\.bash$
+        exclude: ^dev/breeze/autocomplete/.*$
       - id: stylelint
         name: stylelint
         entry: "stylelint"

.rat-excludes

@@ -114,3 +114,6 @@ chart/values_schema.schema.json
 
 # A simplistic Robots.txt
 airflow/www/static/robots.txt
+
+# Generated autocomplete files
+dev/breeze/autocomplete/*

dev/breeze/autocomplete/Breeze2-complete-bash.sh

@@ -0,0 +1,28 @@
+_Breeze2_completion() {
+    local IFS=$'\n'
+    local response
+
+    response=$(env COMP_WORDS="${COMP_WORDS[*]}" COMP_CWORD=$COMP_CWORD _BREEZE2_COMPLETE=bash_complete $1)
+
+    for completion in $response; do
+        IFS=',' read type value <<< "$completion"
+
+        if [[ $type == 'dir' ]]; then
+            COMPREPLY=()
+            compopt -o dirnames
+        elif [[ $type == 'file' ]]; then
+            COMPREPLY=()
+            compopt -o default
+        elif [[ $type == 'plain' ]]; then
+            COMPREPLY+=($value)
+        fi
+    done
+
+    return 0
+}
+
+_Breeze2_completion_setup() {
+    complete -o nosort -F _Breeze2_completion Breeze2
+}
+
+_Breeze2_completion_setup;

dev/breeze/autocomplete/Breeze2-complete-fish.sh

@@ -0,0 +1,21 @@
+function _Breeze2_completion;
+    set -l response;
+
+    for value in (env _BREEZE2_COMPLETE=fish_complete COMP_WORDS=(commandline -cp) COMP_CWORD=(commandline -t) Breeze2);
+        set response $response $value;
+    end;
+
+    for completion in $response;
+        set -l metadata (string split "," $completion);
+
+        if test $metadata[1] = "dir";
+            __fish_complete_directories $metadata[2];
+        else if test $metadata[1] = "file";
+            __fish_complete_path $metadata[2];
+        else if test $metadata[1] = "plain";
+            echo $metadata[2];
+        end;
+    end;
+end;
+
+complete --no-files --command Breeze2 --arguments "(_Breeze2_completion)";

dev/breeze/autocomplete/Breeze2-complete-zsh.sh

@@ -0,0 +1,34 @@
+#compdef Breeze2
+
+_Breeze2_completion() {
+    local -a completions
+    local -a completions_with_descriptions
+    local -a response
+    (( ! $+commands[Breeze2] )) && return 1
+
+    response=("${(@f)$(env COMP_WORDS="${words[*]}" COMP_CWORD=$((CURRENT-1)) _BREEZE2_COMPLETE=zsh_complete Breeze2)}")
+
+    for type key descr in ${response}; do
+        if [[ "$type" == "plain" ]]; then
+            if [[ "$descr" == "_" ]]; then
+                completions+=("$key")
+            else
+                completions_with_descriptions+=("$key":"$descr")
+            fi
+        elif [[ "$type" == "dir" ]]; then
+            _path_files -/
+        elif [[ "$type" == "file" ]]; then
+            _path_files -f
+        fi
+    done
+
+    if [ -n "$completions_with_descriptions" ]; then
+        _describe -V unsorted completions_with_descriptions -U
+    fi
+
+    if [ -n "$completions" ]; then
+        compadd -U -V unsorted -a completions
+    fi
+}
+
+compdef _Breeze2_completion Breeze2;

dev/breeze/setup.cfg

@@ -59,7 +59,6 @@ install_requires =
     pytest-xdist
     rich
     rich_click
-    click_completion
     requests
     psutil
     inputimeout

dev/breeze/src/airflow_breeze/breeze.py

@@ -15,13 +15,19 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import os
+import shutil
 import subprocess
 import sys
 from pathlib import Path
 from typing import Optional, Tuple
 
-import click_completion
-import rich_click as click
+try:
+    import rich_click as click
+except ImportError:
+    # We handle import errors so that click autocomplete works
+    import click  # type: ignore[no-redef]
+
 from click import ClickException
 
 from airflow_breeze.cache import delete_cache, touch_cache_file, write_to_cache_file
@@ -56,15 +62,12 @@
 from airflow_breeze.utils.run_utils import check_package_installed, run_command
 from airflow_breeze.visuals import ASCIIART, ASCIIART_STYLE
 
-AIRFLOW_SOURCES_DIR = Path(__file__).resolve().parent.parent.parent.parent.parent
+AIRFLOW_SOURCES_DIR = Path(__file__).resolve().parent.parent.parent.parent.parent.absolute()
 
 NAME = "Breeze2"
 VERSION = "0.0.1"
 
 
-click_completion.init()
-
-
 @click.group()
 def main():
     find_airflow_sources_root()
@@ -491,66 +494,126 @@ def start_airflow(verbose: bool):
     raise ClickException("\nPlease implement entering breeze.py\n")
 
 
-def write_to_shell(command_to_execute: str, script_path: str, breeze_comment: str):
+BREEZE_COMMENT = "Added by Updated Airflow Breeze autocomplete setup"
+START_LINE = f"# START: {BREEZE_COMMENT}\n"
+END_LINE = f"# END: {BREEZE_COMMENT}\n"
+
+
+def remove_autogenerated_code(script_path: str):
+    lines = Path(script_path).read_text().splitlines(keepends=True)
+    new_lines = []
+    pass_through = True
+    for line in lines:
+        if line == START_LINE:
+            pass_through = False
+            continue
+        if line.startswith(END_LINE):
+            pass_through = True
+            continue
+        if pass_through:
+            new_lines.append(line)
+    Path(script_path).write_text("".join(new_lines))
+
+
+def backup(script_path_file: Path):
+    shutil.copy(str(script_path_file), str(script_path_file) + ".bak")
+
+
+def write_to_shell(command_to_execute: str, script_path: str, force_setup: bool) -> bool:
     skip_check = False
     script_path_file = Path(script_path)
     if not script_path_file.exists():
         skip_check = True
     if not skip_check:
-        with open(script_path) as script_file:
-            if breeze_comment in script_file.read():
-                click.echo("Autocompletion is already setup. Skipping")
-                click.echo(f"Please exit and re-enter your shell or run: \'source {script_path}\'")
-                sys.exit()
-    click.echo(f"This will modify the {script_path} file")
-    with open(script_path, 'a') as script_file:
-        script_file.write(f"\n# START: {breeze_comment}\n")
-        script_file.write(f"{command_to_execute}\n")
-        script_file.write(f"# END: {breeze_comment}\n")
-        click.echo(f"Please exit and re-enter your shell or run: \'source {script_path}\'")
+        if BREEZE_COMMENT in script_path_file.read_text():
+            if not force_setup:
+                console.print(
+                    "\n[yellow]Autocompletion is already setup. Skipping. "
+                    "You can force autocomplete installation by adding --force-setup[/]\n"
+                )
+                return False
+            else:
+                backup(script_path_file)
+                remove_autogenerated_code(script_path)
+    console.print(f"\nModifying the {script_path} file!\n")
+    console.print(f"\nCopy of the file is held in {script_path}.bak !\n")
+    backup(script_path_file)
+    text = script_path_file.read_text()
+    script_path_file.write_text(
+        text + ("\n" if not text.endswith("\n") else "") + START_LINE + command_to_execute + "\n" + END_LINE
+    )
+    console.print(
+        "\n[yellow]Please exit and re-enter your shell or run:[/]\n\n" f"   `source {script_path}`\n"
+    )
+    return True
 
 
+@option_verbose
+@click.option(
+    '-f',
+    '--force-setup',
+    is_flag=True,
+    help='Force autocomplete setup even if already setup before (overrides the setup).',
+)
 @main.command(name='setup-autocomplete')
-def setup_autocomplete():
+def setup_autocomplete(verbose: bool, force_setup: bool):
     """
     Enables autocompletion of Breeze2 commands.
-    Functionality: By default the generated shell scripts will be available in ./dev/breeze/autocomplete/ path
-    Depending on the shell type in the machine we have to link it to the corresponding file
     """
-    global NAME
-    breeze_comment = "Added by Updated Airflow Breeze autocomplete setup"
+
     # Determine if the shell is bash/zsh/powershell. It helps to build the autocomplete path
-    shell = click_completion.get_auto_shell()
-    click.echo(f"Installing {shell} completion for local user")
-    extra_env = {'_CLICK_COMPLETION_COMMAND_CASE_INSENSITIVE_COMPLETE': 'ON'}
-    autocomplete_path = Path(AIRFLOW_SOURCES_DIR) / ".build/autocomplete" / f"{NAME}-complete.{shell}"
-    shell, path = click_completion.core.install(
-        shell=shell, prog_name=NAME, path=autocomplete_path, append=False, extra_env=extra_env
+    detected_shell = os.environ.get('SHELL')
+    detected_shell = None if detected_shell is None else detected_shell.split(os.sep)[-1]
+    if detected_shell not in ['bash', 'zsh', 'fish']:
+        console.print(f"\n[red] The shell {detected_shell} is not supported for autocomplete![/]\n")
+        sys.exit(1)
+    console.print(f"Installing {detected_shell} completion for local user")
+    autocomplete_path = (
+        Path(AIRFLOW_SOURCES_DIR) / "dev" / "breeze" / "autocomplete" / f"{NAME}-complete-{detected_shell}.sh"
+    )
+    console.print(f"[bright_blue]Activation command script is available here: {autocomplete_path}[/]\n")
+    console.print(
+        f"[yellow]We need to add above script to your {detected_shell} profile and "
+        "install 'click' package in your default python installation destination.[/]\n"
     )
-    click.echo(f"Activation command scripts are created in this autocompletion path: {autocomplete_path}")
-    if click.confirm(f"Do you want to add the above autocompletion scripts to your {shell} profile?"):
-        if shell == 'bash':
-            script_path = Path('~').expanduser() / '.bash_completion'
+    updated = False
+    if click.confirm("Should we proceed ?"):
+        if detected_shell == 'bash':
+            script_path = str(Path('~').expanduser() / '.bash_completion')
             command_to_execute = f"source {autocomplete_path}"
-            write_to_shell(command_to_execute, script_path, breeze_comment)
-        elif shell == 'zsh':
-            script_path = Path('~').expanduser() / '.zshrc'
+            updated = write_to_shell(command_to_execute, script_path, force_setup)
+        elif detected_shell == 'zsh':
+            script_path = str(Path('~').expanduser() / '.zshrc')
             command_to_execute = f"source {autocomplete_path}"
-            write_to_shell(command_to_execute, script_path, breeze_comment)
-        elif shell == 'fish':
+            updated = write_to_shell(command_to_execute, script_path, force_setup)
+        elif detected_shell == 'fish':
             # Include steps for fish shell
-            script_path = Path('~').expanduser() / f'.config/fish/completions/{NAME}.fish'
-            with open(path) as source_file, open(script_path, 'w') as destination_file:
-                for line in source_file:
-                    destination_file.write(line)
+            script_path = str(Path('~').expanduser() / f'.config/fish/completions/{NAME}.fish')
+            if os.path.exists(script_path) and not force_setup:
+                console.print(
+                    "\n[yellow]Autocompletion is already setup. Skipping. "
+                    "You can force autocomplete installation by adding --force-setup[/]\n"
+                )
+            else:
+                with open(autocomplete_path) as source_file, open(script_path, 'w') as destination_file:
+                    for line in source_file:
+                        destination_file.write(line)
+                updated = True
         else:
             # Include steps for powershell
             subprocess.check_call(['powershell', 'Set-ExecutionPolicy Unrestricted -Scope CurrentUser'])
-            script_path = subprocess.check_output(['powershell', '-NoProfile', 'echo $profile']).strip()
+            script_path = (
+                subprocess.check_output(['powershell', '-NoProfile', 'echo $profile']).decode("utf-8").strip()
+            )
             command_to_execute = f". {autocomplete_path}"
-            write_to_shell(command_to_execute, script_path.decode("utf-8"), breeze_comment)
+            write_to_shell(command_to_execute, script_path, force_setup)
+        if updated:
+            run_command(['pip', 'install', '--upgrade', 'click'], verbose=True, check=False)
     else:
-        click.echo(f"Link for manually adding the autocompletion script to {shell} profile")
+        console.print(
+            "\nPlease follow the https://click.palletsprojects.com/en/8.1.x/shell-completion/ "
+            "to setup autocompletion for breeze manually if you want to use it.\n"
+        )
 
 
 @main.command(name='config')

dev/breeze/src/airflow_breeze/utils/run_utils.py

@@ -27,9 +27,6 @@
 from pathlib import Path
 from typing import Dict, List, Mapping, Optional
 
-import psutil
-import requests
-
 from airflow_breeze.cache import update_md5checksum_in_cache
 from airflow_breeze.console import console
 from airflow_breeze.global_constants import FILES_FOR_REBUILD_CHECK
@@ -102,6 +99,9 @@ def check_package_installed(package_name: str) -> bool:
 
 
 def get_filesystem_type(filepath):
+    # We import it locally so that click autocomplete works
+    import psutil
+
     root_type = "unknown"
     for part in psutil.disk_partitions():
         if part.mountpoint == '/':
@@ -210,6 +210,9 @@ def fix_group_permissions():
 
 
 def get_latest_sha(repo: str, branch: str):
+    # We import it locally so that click autocomplete works
+    import requests
+
     gh_url = f"https://api.github.com/repos/{repo}/commits/{branch}"
     headers_dict = {"Accept": "application/vnd.github.VERSION.sha"}
     resp = requests.get(gh_url, headers=headers_dict)