Skip to content

Commit

Permalink
use pexpect for tty emulation
Browse files Browse the repository at this point in the history
  • Loading branch information
mjurbanski-reef committed Mar 28, 2024
1 parent 69dd411 commit 5af9ef3
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 31 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
67 changes: 39 additions & 28 deletions b2/_internal/_cli/autocomplete_install.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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')
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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."""


Expand Down
4 changes: 2 additions & 2 deletions b2/_internal/console_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down

0 comments on commit 5af9ef3

Please sign in to comment.