diff --git a/django_typer/management/commands/shellcompletion.py b/django_typer/management/commands/shellcompletion.py index c6225d1..4f9e02f 100644 --- a/django_typer/management/commands/shellcompletion.py +++ b/django_typer/management/commands/shellcompletion.py @@ -26,9 +26,11 @@ import inspect import io import os +import re import sys import typing as t from functools import cached_property +from importlib.resources import files from pathlib import Path from click.parser import split_arg_string @@ -38,6 +40,7 @@ get_completion_class, ) from django.core.management import CommandError, ManagementUtility +from django.template import Context, Engine from django.utils.module_loading import import_string from django.utils.translation import gettext from django.utils.translation import gettext_lazy as _ @@ -48,7 +51,8 @@ completion_init, # pyright: ignore[reportPrivateImportUsage] ) -from django_typer.management import TyperCommand, command, get_command +from django_typer.management import TyperCommand, command, get_command, initialize +from django_typer.types import COMMON_PANEL from django_typer.utils import get_usage_script DETECTED_SHELL = None @@ -94,15 +98,38 @@ class Command(TyperCommand): suppressed_base_arguments = { "version", "skip_checks", - "no_color", - "force_color", "verbosity", } + no_color: t.Optional[bool] = None # type: ignore + _shell: Shells COMPLETE_VAR = "_COMPLETE_INSTRUCTION" + ANSI_ESCAPE_RE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") + + SHELL_COLOR_SUPPORT = { + Shells.bash: False, + Shells.zsh: True, + Shells.fish: True, + Shells.powershell: True, + Shells.pwsh: True, + } + + @cached_property + def template_engine(self): + """ + Django template engine that will find and render completer script templates. + """ + return Engine( + dirs=[str(files("django_typer").joinpath("scripts"))], + libraries={ + "default": "django.template.defaulttags", + "filter": "django.template.defaultfilters", + }, + ) + @cached_property def manage_script(self) -> t.Union[str, Path]: """ @@ -198,66 +225,43 @@ def patch_script( fallback = f" --fallback {fallback}" if fallback else "" - def replace(s: str, old: str, new: str, occurrences: t.List[int]) -> str: - """ - :param s: The string to modify - :param old: The string to replace - :param new: The string to replace with - :param occurrences: A list of occurrences of the old string to replace with the - new string, where the occurrence number is the zero-based count of the old - strings appearance in the string when counting from the front. - """ - count = 0 - result = "" - start = 0 - - for end in range(len(s)): - if s[start : end + 1].endswith(old): - if count in occurrences: - result += f"{s[start:end+1-len(old)]}{new}" - start = end + 1 - else: - result += s[start : end + 1] - start = end + 1 - count += 1 - - result += s[start:] - return result - - if shell is Shells.bash: - typer_scripts._completion_scripts[Shells.bash.value] = replace( - typer_scripts.COMPLETION_SCRIPT_BASH, - "$1", - f"$1 {DJANGO_COMMAND} complete", - [0], - ) - elif shell is Shells.zsh: - typer_scripts._completion_scripts[Shells.zsh.value] = replace( - typer_scripts.COMPLETION_SCRIPT_ZSH, - "%(prog_name)s", - f"${{words[0,1]}} {DJANGO_COMMAND} complete", - [1], - ) - elif shell is Shells.fish: - typer_scripts._completion_scripts[Shells.fish.value] = replace( - typer_scripts.COMPLETION_SCRIPT_FISH, - "%(prog_name)s", - f"{self.manage_script} {DJANGO_COMMAND} complete", - [1, 2], + def render(shell: Shells, **context) -> str: + template = self.template_engine.get_template(f"{shell.value}.tmpl") + return template.render(Context(context)) + + context = { + "manage_script": self.manage_script, + "django_command": DJANGO_COMMAND, + "color": "--no-color" + if self.no_color + else "--force-color" + if self.force_color + else "", + } + if shell in [Shells.powershell, Shells.pwsh]: + typer_scripts._completion_scripts[shell.value] = render( + Shells.powershell, **context ) + elif shell in [Shells.bash, Shells.zsh, Shells.fish]: + typer_scripts._completion_scripts[shell.value] = render(shell, **context) else: - assert shell in [ - Shells.pwsh, - Shells.powershell, - ], gettext("Unsupported shell: {shell}").format(shell=shell.value) - script = replace( - typer_scripts.COMPLETION_SCRIPT_POWER_SHELL, - "%(prog_name)s", - f"{self.manage_script} {DJANGO_COMMAND} complete", - [0], + raise NotImplementedError( + gettext("Unsupported shell: {shell}").format(shell=shell.value) ) - typer_scripts._completion_scripts[Shells.powershell.value] = script - typer_scripts._completion_scripts[Shells.pwsh.value] = script + + @initialize() + def init( + self, + no_color: t.Annotated[ + t.Optional[bool], + Option( + "--no-color", + help=t.cast(str, _("Filter color codes out of completion text.")), + rich_help_panel=COMMON_PANEL, + ), + ] = None, + ): + self.no_color = no_color # pyright: ignore[reportAttributeAccessIssue] @command( help=t.cast(str, _("Install autocompletion for the current or given shell.")) @@ -312,6 +316,12 @@ def install( from typer._completion_shared import install self.shell = shell # type: ignore + assert self.shell + self.no_color = ( # pyright: ignore[reportIncompatibleVariableOverride] + not self.SHELL_COLOR_SUPPORT.get(self.shell, False) + if self.no_color is None + else self.no_color + ) self.patch_script(fallback=fallback) install_path = install( shell=self.shell.value, @@ -599,6 +609,14 @@ def get_completion() -> None: ) call_fallback(fallback) + def strip_color(text: str) -> str: + """ + Strip ANSI color codes from a string. + """ + if self.no_color: + return self.ANSI_ESCAPE_RE.sub("", text) + return text + buffer = io.StringIO() try: with contextlib.redirect_stdout(buffer): @@ -608,11 +626,11 @@ def get_completion() -> None: get_completion() # leave this in, incase the interface changes to not exit - return buffer.getvalue() # pragma: no cover + return strip_color(buffer.getvalue()) # pragma: no cover except SystemExit: completion_str = buffer.getvalue() if completion_str: - return completion_str + return strip_color(completion_str) if cmd_str: return "" raise diff --git a/django_typer/scripts/bash.tmpl b/django_typer/scripts/bash.tmpl new file mode 100644 index 0000000..02ad3da --- /dev/null +++ b/django_typer/scripts/bash.tmpl @@ -0,0 +1,10 @@ +%(complete_func)s() { + local IFS=$' +' + COMPREPLY=( $( env COMP_WORDS="${COMP_WORDS[*]}" \ + COMP_CWORD=$COMP_CWORD \ + %(autocomplete_var)s=complete_bash $1 {{ django_command }} {{ color }} complete ) ) + return 0 +} + +complete -o default -F %(complete_func)s %(prog_name)s diff --git a/django_typer/scripts/fish.tmpl b/django_typer/scripts/fish.tmpl new file mode 100644 index 0000000..7e8e846 --- /dev/null +++ b/django_typer/scripts/fish.tmpl @@ -0,0 +1 @@ +complete --command %(prog_name)s --no-files --arguments "(env %(autocomplete_var)s=complete_fish _TYPER_COMPLETE_FISH_ACTION=get-args _TYPER_COMPLETE_ARGS=(commandline -cp) {{ manage_script }} {{ django_command }} {{ color }} complete)" --condition "env %(autocomplete_var)s=complete_fish _TYPER_COMPLETE_FISH_ACTION=is-args _TYPER_COMPLETE_ARGS=(commandline -cp) {{ manage_script }} {{ django_command }} {{ color }} complete" diff --git a/django_typer/scripts/powershell.tmpl b/django_typer/scripts/powershell.tmpl new file mode 100644 index 0000000..b835f51 --- /dev/null +++ b/django_typer/scripts/powershell.tmpl @@ -0,0 +1,19 @@ +Import-Module PSReadLine +Set-PSReadLineKeyHandler -Chord Tab -Function MenuComplete +$scriptblock = { + param($wordToComplete, $commandAst, $cursorPosition) + $Env:%(autocomplete_var)s = "complete_powershell" + $Env:_TYPER_COMPLETE_ARGS = $commandAst.ToString() + $Env:_TYPER_COMPLETE_WORD_TO_COMPLETE = $wordToComplete + {{ manage_script }} {{ django_command }} {{ color }} complete | ForEach-Object { + $commandArray = $_ -Split ":::" + $command = $commandArray[0] + $helpString = $commandArray[1] + [System.Management.Automation.CompletionResult]::new( + $command, $command, 'ParameterValue', $helpString) + } + $Env:%(autocomplete_var)s = "" + $Env:_TYPER_COMPLETE_ARGS = "" + $Env:_TYPER_COMPLETE_WORD_TO_COMPLETE = "" +} +Register-ArgumentCompleter -Native -CommandName %(prog_name)s -ScriptBlock $scriptblock diff --git a/django_typer/scripts/zsh.tmpl b/django_typer/scripts/zsh.tmpl new file mode 100644 index 0000000..8b02048 --- /dev/null +++ b/django_typer/scripts/zsh.tmpl @@ -0,0 +1,7 @@ +#compdef %(prog_name)s + +%(complete_func)s() { + eval $(env _TYPER_COMPLETE_ARGS="${words[1,$CURRENT]}" %(autocomplete_var)s=complete_zsh ${words[0,1]} {{ django_command }} {{ color }} complete) +} + +compdef %(complete_func)s %(prog_name)s diff --git a/doc/source/changelog.rst b/doc/source/changelog.rst index 03146e3..d50014b 100644 --- a/doc/source/changelog.rst +++ b/doc/source/changelog.rst @@ -6,6 +6,7 @@ v3.0.0 (202X-XX-XX) =================== * Fixed `Typer-style interface throws an assertion when no callback is present on a subgroup. `_ +* Implemented `ANSI color control sequences should optionally be scrubbed from shell completions `_ * Fixed `supressed_base_arguments are still present in the Context `_ * Implemented `Add showcase of commands using django-typer to docs `_ * Implemented `Add a @finalize decorator for functions to collect/operate on subroutine results. `_ diff --git a/tests/test_parser_completers.py b/tests/test_parser_completers.py index b956f89..99ab7ac 100644 --- a/tests/test_parser_completers.py +++ b/tests/test_parser_completers.py @@ -673,6 +673,7 @@ def test_uuid_field(self): with contextlib.redirect_stdout(result): call_command( "shellcompletion", + "--no-color", "complete", "model_fields test --uuid 123456&78-^56785678f234---A", ) @@ -940,7 +941,9 @@ def test_decimal_field(self): def test_option_complete(self): result = StringIO() with contextlib.redirect_stdout(result): - call_command("shellcompletion", "complete", "model_fields test ") + call_command( + "shellcompletion", "--no-color", "complete", "model_fields test " + ) result = result.getvalue() self.assertTrue("--char" in result) self.assertTrue("--ichar" in result)