From 5af9ef3898a2bce8492b5e9488e482f8a13f790e Mon Sep 17 00:00:00 2001
From: Maciej Urbanski <maciej.rooter.urbanski@reef.pl>
Date: Thu, 28 Mar 2024 11:02:24 +0100
Subject: [PATCH] use pexpect for tty emulation

---
 .github/workflows/ci.yml                  |  2 +-
 b2/_internal/_cli/autocomplete_install.py | 67 +++++++++++++----------
 b2/_internal/console_tool.py              |  4 +-
 3 files changed, 42 insertions(+), 31 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ec9a0b312..7aaac3076 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -103,7 +103,7 @@ jobs:
       - name: Install test binary dependencies (macOS)
         if: startsWith(matrix.os, 'macos')
         run: |
-          brew install fish expect
+          brew install fish
       - name: Install dependencies
         run: python -m pip install --upgrade nox pdm
       - name: Run unit tests
diff --git a/b2/_internal/_cli/autocomplete_install.py b/b2/_internal/_cli/autocomplete_install.py
index ade161303..452137c35 100644
--- a/b2/_internal/_cli/autocomplete_install.py
+++ b/b2/_internal/_cli/autocomplete_install.py
@@ -10,14 +10,16 @@
 from __future__ import annotations
 
 import abc
+import io
 import logging
 import os
-import platform
 import re
 import shutil
+import signal
 import subprocess
 import textwrap
 from datetime import datetime
+from importlib.util import find_spec
 from pathlib import Path
 from shlex import quote
 
@@ -169,8 +171,16 @@ def is_enabled(self) -> bool:
         if not rc_path.exists():
             # if zshrc is missing `zshrc -i` may hang on creation wizard when emulating tty
             rc_path.touch(mode=0o750)
+        emulate_tty = os.isatty(0)  # is True under GHA or pytest-xdist
+        if emulate_tty and not find_spec('pexpect'):
+            emulate_tty = False
+            logger.warning(
+                "pexpect is recommended for Zsh shell autocomplete installation check without tty. "
+                "You can install it via `pip install pexpect`."
+            )
+
         cmd = [self.shell_exec, '-i', '-c', f'[[ -v _comps[{quote(self.prog)}] ]]']
-        return _silent_success_run(cmd) if os.isatty(0) else _silent_success_run(_tty_cmd_wrap(cmd))
+        return _silent_success_run(cmd) if emulate_tty else _silent_success_run_with_pty(cmd)
 
 
 @SHELL_REGISTRY.register('fish')
@@ -221,32 +231,37 @@ def _silent_success_run(cmd: list[str], timeout: int | None = 60, env: dict | No
         stdout, stderr = p.communicate(timeout=1)
         logger.warning("Command %r timed out, stdout: %r, stderr: %r", cmd, stdout, stderr)
     else:
-        if p.returncode != 0:
-            logger.debug(
-                "Command %r exited with code %d, stdout: %r, stderr: %r", cmd, p.returncode, stdout,
-                stderr
-            )
+        logger.debug(
+            "Command %r exited with code %r, stdout: %r, stderr: %r", cmd, p.returncode, stdout,
+            stderr
+        )
     return p.returncode == 0
 
 
-def _tty_cmd_wrap(cmd: list[str]) -> list[str]:
+def _silent_success_run_with_pty(cmd: list[str], timeout: int = 5, env: dict | None = None) -> bool:
     """
-    Add tty emulation to the command.
-
-    Inspired by tty emulation for different platforms found here:
-     https://github.com/Yuri6037/Action-FakeTTY/tree/master/script
+    Run a command with emulated terminal and return whether it succeeded.
     """
-    system_name = platform.system().lower()
-    if system_name == 'darwin':
-        if not shutil.which('unbuffer'):
-            raise CLIError(
-                "unbuffer is required for macOS when running without a tty. "
-                "You can install it via `brew install expect`."
-            )
-        return ['unbuffer', *cmd]
-    elif system_name == 'windows':
-        return ["Invoke-Expression", shlex_join(cmd)]
-    return ['script', '-q', '-e', '-c', shlex_join(cmd), '/dev/null']
+    import pexpect
+
+    command_str = shlex_join(cmd)
+
+    child = pexpect.spawn(command_str, timeout=timeout, env=env)
+    output = io.BytesIO()
+    try:
+        child.logfile_read = output
+        child.expect(pexpect.EOF)
+    except pexpect.TIMEOUT:
+        logger.warning("Command %r timed out, output: %r", cmd, output.getvalue())
+        child.kill(signal.SIGKILL)
+        return False
+    finally:
+        child.close()
+
+    logger.debug(
+        "Command %r exited with code %r, output: %r", cmd, child.exitstatus, output.getvalue()
+    )
+    return child.exitstatus == 0
 
 
 def add_or_update_shell_section(
@@ -278,11 +293,7 @@ def add_or_update_shell_section(
     path.write_text(file_content)
 
 
-class CLIError(Exception):
-    """Base exception for CLI errors."""
-
-
-class AutocompleteInstallError(CLIError):
+class AutocompleteInstallError(Exception):
     """Exception raised when autocomplete installation fails."""
 
 
diff --git a/b2/_internal/console_tool.py b/b2/_internal/console_tool.py
index 68154f485..e60618f3b 100644
--- a/b2/_internal/console_tool.py
+++ b/b2/_internal/console_tool.py
@@ -127,7 +127,7 @@
 from b2._internal._cli.argcompleters import bucket_name_completer, file_name_completer
 from b2._internal._cli.autocomplete_install import (
     SUPPORTED_SHELLS,
-    CLIError,
+    AutocompleteInstallError,
     autocomplete_install,
 )
 from b2._internal._cli.b2api import _get_b2api_for_profile
@@ -4165,7 +4165,7 @@ def _run(self, args):
 
         try:
             autocomplete_install(self.console_tool.b2_binary_name, shell=shell)
-        except CLIError as e:
+        except AutocompleteInstallError as e:
             raise CommandError(str(e)) from e
         self._print(f'Autocomplete successfully installed for {shell}.')
         self._print(