fix #144

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

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

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"

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

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

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. <https://github.com/django-commons/django-typer/issues/145>`_
+* Implemented `ANSI color control sequences should optionally be scrubbed from shell completions <https://github.com/django-commons/django-typer/issues/144>`_
 * Fixed `supressed_base_arguments are still present in the Context <https://github.com/django-commons/django-typer/issues/143>`_
 * Implemented `Add showcase of commands using django-typer to docs <https://github.com/django-commons/django-typer/issues/142>`_
 * Implemented `Add a @finalize decorator for functions to collect/operate on subroutine results. <https://github.com/django-commons/django-typer/issues/140>`_

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)