From 393d19914b44107bae9aa1488ddddd7d80ee2d20 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Fri, 1 Apr 2022 20:35:15 +0200 Subject: [PATCH] Fix Breeze2 autocomplete 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 + .../autocomplete/Breeze2-complete-bash.sh | 28 ++++ .../autocomplete/Breeze2-complete-fish.sh | 21 +++ .../autocomplete/Breeze2-complete-zsh.sh | 34 ++++ dev/breeze/setup.cfg | 1 - dev/breeze/src/airflow_breeze/breeze.py | 153 ++++++++++++------ .../src/airflow_breeze/utils/run_utils.py | 9 +- 8 files changed, 202 insertions(+), 50 deletions(-) create mode 100644 dev/breeze/autocomplete/Breeze2-complete-bash.sh create mode 100644 dev/breeze/autocomplete/Breeze2-complete-fish.sh create mode 100644 dev/breeze/autocomplete/Breeze2-complete-zsh.sh diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1c8a1cdbf7489..03f5076161004 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 94a97319f2a3c..53d477664df53 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 0000000000000..a9ba398985c0c --- /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 0000000000000..7dea0bc38d79d --- /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 0000000000000..11092e75f3fbe --- /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 18ffcdc9a2b10..3408a5523b394 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 1a93a4f1a2eab..e6ea13fcd44ef 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.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 f48d16e54011e..8553e3a8f28c1 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 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)