powershell completion tests almost working #156

bckohan · bckohan · commit e5502f426a82 · 2025-01-04T04:13:29.000-08:00
diff --git a/django_typer/templates/shell_complete/powershell.ps1 b/django_typer/templates/shell_complete/powershell.ps1
@@ -20,7 +20,14 @@ $scriptblock = {
         $pythonPathOption = "--pythonpath=$($matches[1])"
     }
 
-    {{ manage_script_name }} {{ django_command }} $settingsOption $pythonPathOption --shell {{ shell }} {{ color }} complete "$($commandText)" | ForEach-Object {
+    $results = {{ manage_script_name }} {{ django_command }} $settingsOption $pythonPathOption --shell {{ shell }} {{ color }} complete "$($commandText)"
+
+    if ($results.Count -eq 0) {
+        # avoid default path completion
+        return $null
+    }
+
+    $results | ForEach-Object {
         $commandArray = $_ -Split ":::"
         $type = $commandArray[0]
         $value = $commandArray[1]
diff --git a/django_typer/utils.py b/django_typer/utils.py
@@ -236,36 +236,32 @@ def get_win_shell() -> str:
     from shellingham import ShellDetectionFailure
 
     assert platform.system() == "Windows"
-    pwsh = shutil.which("pwsh")
-    powershell = shutil.which("powershell")
-    if pwsh and not powershell:
-        return "pwsh"
-    elif powershell and not pwsh:
-        return "powershell"
-    try:
-        ps_command = """
-        $parent = Get-CimInstance -Query "SELECT * FROM Win32_Process WHERE ProcessId = {pid}";
-        $parentPid = $parent.ParentProcessId;
-        $parentInfo = Get-CimInstance -Query "SELECT * FROM Win32_Process WHERE ProcessId = $parentPid";
-        $parentInfo | Select-Object Name, ProcessId | ConvertTo-Json -Depth 1
-        """
-        pid = os.getpid()
-        while True:
-            result = subprocess.run(
-                ["pwsh", "-NoProfile", "-Command", ps_command.format(pid=pid)],
-                capture_output=True,
-                text=True,
-            ).stdout.strip()
-            if not result:
-                break
-            process = json.loads(result)
-            if "pwsh" in process.get("Name", ""):
-                return "pwsh"
-            elif "powershell" in process.get("Name", ""):
-                return "powershell"
-            pid = process["ProcessId"]
-
-        raise ShellDetectionFailure("Unable to detect windows shell")
-
-    except Exception as e:
-        raise ShellDetectionFailure("Unable to detect windows shell") from e
+    pwsh = shutil.which("pwsh") or shutil.which("powershell")
+    if pwsh:
+        try:
+            ps_command = """
+            $parent = Get-CimInstance -Query "SELECT * FROM Win32_Process WHERE ProcessId = {pid}";
+            $parentPid = $parent.ParentProcessId;
+            $parentInfo = Get-CimInstance -Query "SELECT * FROM Win32_Process WHERE ProcessId = $parentPid";
+            $parentInfo | Select-Object Name, ProcessId | ConvertTo-Json -Depth 1
+            """
+            pid = os.getpid()
+            while True:
+                result = subprocess.run(
+                    [pwsh, "-NoProfile", "-Command", ps_command.format(pid=pid)],
+                    capture_output=True,
+                    text=True,
+                ).stdout.strip()
+                if not result:
+                    break
+                process = json.loads(result)
+                if "pwsh" in process.get("Name", ""):
+                    return "pwsh"
+                elif "powershell" in process.get("Name", ""):
+                    return "powershell"
+                pid = process["ProcessId"]
+
+        except Exception as e:
+            raise ShellDetectionFailure("Unable to detect windows shell") from e
+
+    raise ShellDetectionFailure("Unable to detect windows shell")
diff --git a/tests/shellcompletion/__init__.py b/tests/shellcompletion/__init__.py
@@ -1,19 +1,17 @@
-import fcntl
 import os
-import pty
 import re
 import select
 import struct
 import subprocess
 import sys
-import termios
 import time
 import typing as t
 from pathlib import Path
 import pytest
 from functools import cached_property
 import re
 import subprocess
+import platform
 
 from shellingham import detect_shell
 
@@ -67,11 +65,9 @@ class _DefaultCompleteTestCase(with_typehint(TestCase)):
     manage_script = "manage.py"
     launch_script = "./manage.py"
 
-    @property
-    def interactive_opt(self):
-        # currently all supported shells support -i for interactive mode
-        # this includes zsh, bash, fish and powershell
-        return "-i"
+    interactive_opt: t.Optional[str] = None
+
+    environment: t.List[str] = []
 
     @cached_property
     def command(self) -> ShellCompletion:
@@ -116,75 +112,109 @@ def remove(self, script=None):
         self.get_completions("ping")  # just to reinit shell
         self.verify_remove(script=script)
 
-    def set_environment(self, fd):
-        os.write(fd, f"PATH={Path(sys.executable).parent}:$PATH\n".encode())
-        os.write(
-            fd,
-            f"DJANGO_SETTINGS_MODULE=tests.settings.completion\n".encode(),
-        )
+    if platform.system() == "Windows":
 
-    def get_completions(self, *cmds: str, scrub_output=True) -> str:
-        def read(fd):
-            """Function to read from a file descriptor."""
-            return os.read(fd, 1024 * 1024).decode()
-
-        # Create a pseudo-terminal
-        master_fd, slave_fd = pty.openpty()
-
-        # Define window size - width and height
-        os.set_blocking(slave_fd, False)
-        win_size = struct.pack("HHHH", 24, 80, 0, 0)  # 24 rows, 80 columns
-        fcntl.ioctl(slave_fd, termios.TIOCSWINSZ, win_size)
-
-        # env = os.environ.copy()
-        # env["TERM"] = "xterm-256color"
-
-        # Spawn a new shell process
-        shell = self.shell or detect_shell()[0]
-        process = subprocess.Popen(
-            [shell, *([self.interactive_opt] if self.interactive_opt else [])],
-            stdin=slave_fd,
-            stdout=slave_fd,
-            stderr=slave_fd,
-            text=True,
-            # env=env,
-            # preexec_fn=os.setsid,
-        )
-        # Wait for the shell to start and get to the prompt
-        print(read(master_fd))
+        def get_completions(self, *cmds: str, scrub_output=True) -> str:
+            import winpty
+
+            assert self.shell
+
+            pty = winpty.PTY(24, 80)
+
+            def read_all() -> str:
+                output = ""
+                while data := pty.read():
+                    output += data
+                    time.sleep(0.1)
+                return output
+
+            # Start the subprocess
+            pty.spawn(
+                self.shell, *([self.interactive_opt] if self.interactive_opt else [])
+            )
+
+            # Wait for the shell to start and get to the prompt
+            time.sleep(3)
+            read_all()
+
+            for line in self.environment:
+                pty.write(line)
+
+            time.sleep(2)
+            output = read_all() + read_all()
+
+            pty.write(" ".join(cmds))
+            time.sleep(0.1)
+            pty.write("\t")
+
+            time.sleep(2)
+            completion = read_all() + read_all()
 
-        self.set_environment(master_fd)
+            return scrub(completion) if scrub_output else completion
 
-        print(read(master_fd))
-        # Send a command with a tab character for completion
+    else:
 
-        cmd = " ".join(cmds)
-        os.write(master_fd, cmd.encode())
-        time.sleep(0.25)
+        def get_completions(self, *cmds: str, scrub_output=True) -> str:
+            import fcntl
+            import termios
+            import pty
 
-        print(f'"{cmd}"')
-        os.write(master_fd, b"\t\t\t")
+            def read(fd):
+                """Function to read from a file descriptor."""
+                return os.read(fd, 1024 * 1024).decode()
 
-        time.sleep(0.25)
+            # Create a pseudo-terminal
+            master_fd, slave_fd = pty.openpty()
 
-        # Read the output
-        output = read_all_from_fd_with_timeout(master_fd)
+            # Define window size - width and height
+            os.set_blocking(slave_fd, False)
+            win_size = struct.pack("HHHH", 24, 80, 0, 0)  # 24 rows, 80 columns
+            fcntl.ioctl(slave_fd, termios.TIOCSWINSZ, win_size)
 
-        # todo - avoid large output because this can mess things up
-        if "do you wish" in output or "Display all" in output:
-            os.write(master_fd, b"y\n")
+            # env = os.environ.copy()
+            # env["TERM"] = "xterm-256color"
+
+            # Spawn a new shell process
+            shell = self.shell or detect_shell()[0]
+            process = subprocess.Popen(
+                [shell, *([self.interactive_opt] if self.interactive_opt else [])],
+                stdin=slave_fd,
+                stdout=slave_fd,
+                stderr=slave_fd,
+                text=True,
+                # env=env,
+                # preexec_fn=os.setsid,
+            )
+            # Wait for the shell to start and get to the prompt
+            print(read(master_fd))
+
+            for line in self.environment:
+                os.write(master_fd, line.encode())
+
+            print(read(master_fd))
+            # Send a command with a tab character for completion
+
+            cmd = " ".join(cmds)
+            os.write(master_fd, cmd.encode())
             time.sleep(0.25)
+
+            print(f'"{cmd}"')
+            os.write(master_fd, b"\t\t\t")
+
+            time.sleep(0.25)
+
+            # Read the output
             output = read_all_from_fd_with_timeout(master_fd)
 
-        # Clean up
-        os.close(slave_fd)
-        os.close(master_fd)
-        process.terminate()
-        process.wait()
-        # remove bell character which can show up in some terminals where we hit tab
-        if scrub_output:
-            return scrub(output)
-        return output
+            # Clean up
+            os.close(slave_fd)
+            os.close(master_fd)
+            process.terminate()
+            process.wait()
+            # remove bell character which can show up in some terminals where we hit tab
+            if scrub_output:
+                return scrub(output)
+            return output
 
     def run_app_completion(self):
         completions = self.get_completions(self.launch_script, "completion", " ")
@@ -323,18 +353,24 @@ class _InstalledScriptTestCase(_DefaultCompleteTestCase):
     manage_script = "django_manage"
     launch_script = "django_manage"
 
-    def setUp(self):
+    @classmethod
+    def setUpClass(cls):
         lines = []
-        with open(self.MANAGE_SCRIPT_TMPL, "r") as f:
+        with open(cls.MANAGE_SCRIPT_TMPL, "r") as f:
             for line in f.readlines():
                 if line.startswith("#!{{shebang}}"):
                     line = f"#!{sys.executable}\n"
                 lines.append(line)
-        exe = Path(sys.executable).parent / self.manage_script
+        exe = Path(sys.executable).parent / cls.manage_script
         with open(exe, "w") as f:
             for line in lines:
                 f.write(line)
 
         # make the script executable
         os.chmod(exe, os.stat(exe).st_mode | 0o111)
-        super().setUp()
+
+        if platform.system() == "Windows":
+            with open(exe.with_suffix(".cmd"), "w") as f:
+                f.write(f'@echo off{os.linesep}"{sys.executable}" "%~dp0{exe.name}" %*')
+            os.chmod(exe, os.stat(exe.with_suffix(".cmd")).st_mode | 0o111)
+        super().setUpClass()
diff --git a/tests/shellcompletion/test_bash.py b/tests/shellcompletion/test_bash.py
@@ -13,16 +13,13 @@
 class BashTests(_DefaultCompleteTestCase, TestCase):
     shell = "bash"
     directory = Path("~/.bash_completions").expanduser()
+    interactive_opt = "-i"
 
-    def set_environment(self, fd):
-        os.write(
-            fd,
-            f"export DJANGO_SETTINGS_MODULE=tests.settings.completion\n".encode(),
-        )
-        os.write(fd, "source ~/.bashrc\n".encode())
-        activate = Path(sys.executable).absolute().parent / "activate"
-        if activate.is_file():
-            os.write(fd, f"source {activate}\n".encode())
+    environment = [
+        f"export DJANGO_SETTINGS_MODULE=tests.settings.completion{os.linesep}",
+        f"source ~/.bashrc{os.linesep}",
+        f"source {Path(sys.executable).absolute().parent / 'activate'}{os.linesep}",
+    ]
 
     def verify_install(self, script=None):
         if not script:
diff --git a/tests/shellcompletion/test_powershell.py b/tests/shellcompletion/test_powershell.py
diff --git a/tests/shellcompletion/test_zsh.py b/tests/shellcompletion/test_zsh.py
