You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2022/04/03 18:54:11 UTC

[airflow] branch main updated: Fix Breeze2 autocomplete (#22695)

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 4fb929a  Fix Breeze2 autocomplete (#22695)
4fb929a is described below

commit 4fb929a60966e9cecbbb435efd375adcc4fff9d7
Author: Jarek Potiuk <ja...@potiuk.com>
AuthorDate: Sun Apr 3 20:52:48 2022 +0200

    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
---
 .pre-commit-config.yaml                          |   3 +-
 .rat-excludes                                    |   3 +
 dev/breeze/autocomplete/Breeze2-complete-bash.sh |  28 +++++
 dev/breeze/autocomplete/Breeze2-complete-fish.sh |  21 ++++
 dev/breeze/autocomplete/Breeze2-complete-zsh.sh  |  34 +++++
 dev/breeze/setup.cfg                             |   1 -
 dev/breeze/src/airflow_breeze/breeze.py          | 153 ++++++++++++++++-------
 dev/breeze/src/airflow_breeze/utils/run_utils.py |   9 +-
 8 files changed, 202 insertions(+), 50 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1c8a1cd..03f5076 100644
--- a/.pre-commit-config.yaml
+++ b/.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"
diff --git a/.rat-excludes b/.rat-excludes
index 94a9731..53d4776 100644
--- a/.rat-excludes
+++ b/.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/*
diff --git a/dev/breeze/autocomplete/Breeze2-complete-bash.sh b/dev/breeze/autocomplete/Breeze2-complete-bash.sh
new file mode 100644
index 0000000..a9ba398
--- /dev/null
+++ b/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;
diff --git a/dev/breeze/autocomplete/Breeze2-complete-fish.sh b/dev/breeze/autocomplete/Breeze2-complete-fish.sh
new file mode 100644
index 0000000..7dea0bc
--- /dev/null
+++ b/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)";
diff --git a/dev/breeze/autocomplete/Breeze2-complete-zsh.sh b/dev/breeze/autocomplete/Breeze2-complete-zsh.sh
new file mode 100644
index 0000000..11092e7
--- /dev/null
+++ b/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;
diff --git a/dev/breeze/setup.cfg b/dev/breeze/setup.cfg
index 18ffcdc..3408a55 100644
--- a/dev/breeze/setup.cfg
+++ b/dev/breeze/setup.cfg
@@ -59,7 +59,6 @@ install_requires =
     pytest-xdist
     rich
     rich_click
-    click_completion
     requests
     psutil
     inputimeout
diff --git a/dev/breeze/src/airflow_breeze/breeze.py b/dev/breeze/src/airflow_breeze/breeze.py
index 1a93a4f..e6ea13f 100755
--- a/dev/breeze/src/airflow_breeze/breeze.py
+++ b/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.path_utils import (
 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')
diff --git a/dev/breeze/src/airflow_breeze/utils/run_utils.py b/dev/breeze/src/airflow_breeze/utils/run_utils.py
index f48d16e..8553e3a 100644
--- a/dev/breeze/src/airflow_breeze/utils/run_utils.py
+++ b/dev/breeze/src/airflow_breeze/utils/run_utils.py
@@ -27,9 +27,6 @@ from copy import deepcopy
 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)