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)