Fix Breeze2 autocomplete

.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)