Repository: byllyfish/shellous (language: python)

Issue:
Add the Prompt class for working with interactive programs
Add a utility class named Prompt that helps work with an interactive prompt:

It is used with the Runner class like so:

```python
    async with sh(script).set(pty=True).run() as net:
        prompt = Prompt(net, "mininet> ")

        welcome = await prompt.send()
        if welcome.startswith("Error:"):
            raise RuntimeError(welcome)

        # Do stuff with prompt here.
        result = await prompt.send("pingall")
        assert "100%" in result

        await prompt.send("exit")
```


Discussion / hints:

Fix patch:
diff --git a/.cspell.json b/.cspell.json
index 93ef2a3f..bd9cf71d 100644
--- a/.cspell.json
+++ b/.cspell.json
@@ -26,6 +26,7 @@
     ],
     "ignoreRegExpList": [
         "# pylint:.*",
-        "`[a-z]+`"
+        "`[a-z]+`",
+        "__[a-z]+__"
     ]
 }
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index e2d1499d..69f91219 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,7 +21,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-22.04, ubuntu-20.04, macos-12, macos-11, windows-2022 ]
+        os: [ ubuntu-22.04, ubuntu-20.04, macos-13, macos-12, macos-11, windows-2022 ]
         python-version: [ '3.9', '3.10', '3.11', '3.12-dev']
 
     env:
diff --git a/ci/annotate.awk b/ci/annotate.awk
index c0a747e5..55753b20 100755
--- a/ci/annotate.awk
+++ b/ci/annotate.awk
@@ -6,7 +6,7 @@
 #     pytest -v -s --log-cli-level=DEBUG | ./ci/annotate.awk -v step=StepName
 
 /CRITICAL/ { critical_msg[critical_count++] = $0 }
-/[^_]ERROR/ { if (!match($0, /test_logger/)) error_msg[error_count++] = $0 }
+/[^_.]ERROR/ { if (!match($0, /test_logger/)) error_msg[error_count++] = $0 }
 /WARNING/  { warning_msg[warning_count++] = $0 }
 1          # print every line
 
diff --git a/shellous/prompt.py b/shellous/prompt.py
new file mode 100644
index 00000000..7e4d4331
--- /dev/null
+++ b/shellous/prompt.py
@@ -0,0 +1,131 @@
+"Implements the Prompt utility class."
+
+import asyncio
+from typing import Optional
+
+from shellous.harvest import harvest_results
+from shellous.log import LOGGER
+from shellous.runner import Runner
+
+
+class Prompt:
+    """Utility class to help with an interactive prompt session.
+
+    This is an experimental API.
+
+    Example:
+    ```
+    cmd = sh("sh").stdin(sh.CAPTURE).stdout(sh.CAPTURE).stderr(sh.STDOUT)
+
+    async with cmd.env(PS1="??? ") as run:
+        prompt = Prompt(run, "??? ")
+        greeting = await prompt.receive()
+
+        result = await prompt.send("echo hello")
+        assert result == "hello\n"
+
+        prompt.close()
+    ```
+    """
+
+    runner: Runner
+    prompt_bytes: bytes
+    default_timeout: Optional[float]
+
+    def __init__(
+        self,
+        runner: Runner,
+        prompt: str,
+        *,
+        default_timeout: Optional[float] = None,
+    ):
+        assert runner.stdin is not None
+        assert runner.stdout is not None
+
+        self.runner = runner
+        self.prompt_bytes = prompt.encode("utf-8")
+        self.default_timeout = default_timeout
+
+    async def send(
+        self,
+        input_text: str,
+        *,
+        timeout: Optional[float] = None,
+    ) -> str:
+        "Write some input text to stdin, then await the response."
+        stdin = self.runner.stdin
+        assert stdin is not None
+
+        data = input_text.encode("utf-8") + b"\n"
+        LOGGER.debug("Prompt[pid=%s] send: %r", self.runner.pid, data)
+        stdin.write(data)
+
+        # Drain our write to stdin and wait for prompt from stdout.
+        cancelled, (result, _) = await harvest_results(
+            self._read_to_prompt(),
+            stdin.drain(),
+            timeout=timeout or self.default_timeout,
+        )
+        if cancelled:
+            raise asyncio.CancelledError()
+
+        return result
+
+    async def receive(
+        self,
+        *,
+        timeout: Optional[float] = None,
+    ):
+        "Read from stdout up to the next prompt."
+        cancelled, (result,) = await harvest_results(
+            self._read_to_prompt(),
+            timeout=timeout or self.default_timeout,
+        )
+        if cancelled:
+            raise asyncio.CancelledError()
+
+        return result
+
+    def close(self) -> None:
+        "Close stdin to end the prompt session."
+        assert self.runner.stdin is not None
+        self.runner.stdin.close()
+
+    async def _read_to_prompt(self) -> str:
+        "Read all data up to the prompt and return it (after removing prompt)."
+        stdout = self.runner.stdout
+        assert stdout is not None
+
+        buf = await _read_until(stdout, self.prompt_bytes)
+        LOGGER.debug("Prompt[pid=%s] receive: %r", self.runner.pid, buf)
+
+        # Clean up the output to remove the prompt, then return as string.
+        buf = buf.replace(b"\r\n", b"\n")
+        if buf.endswith(self.prompt_bytes):
+            buf = buf[0 : -len(self.prompt_bytes)].rstrip(b"\n")
+
+        return buf.decode("utf-8")
+
+
+async def _read_until(stream: asyncio.StreamReader, separator: bytes) -> bytes:
+    "Read all data until the separator."
+    try:
+        # Most reads can complete without buffering.
+        return await stream.readuntil(separator)
+    except asyncio.IncompleteReadError as ex:
+        return ex.partial
+    except asyncio.LimitOverrunError as ex:
+        # Okay, we have to buffer.
+        buf = bytearray(await stream.read(ex.consumed))
+
+    while True:
+        try:
+            buf.extend(await stream.readuntil(separator))
+        except asyncio.IncompleteReadError as ex:
+            buf.extend(ex.partial)
+        except asyncio.LimitOverrunError as ex:
+            buf.extend(await stream.read(ex.consumed))
+            continue
+        break
+
+    return bytes(buf)


Test patch:
diff --git a/tests/test_prompt.py b/tests/test_prompt.py
new file mode 100644
index 00000000..48d62601
--- /dev/null
+++ b/tests/test_prompt.py
@@ -0,0 +1,274 @@
+"Unit tests for the Prompt class."
+
+import asyncio
+import os
+import platform
+import sys
+
+import pytest
+
+from shellous import cooked, sh
+from shellous.prompt import Prompt
+
+# True if we are running on PyPy.
+_IS_PYPY = platform.python_implementation() == "PyPy"
+
+# True if we're running on alpine linux.
+_IS_ALPINE = os.path.exists("/etc/alpine-release")
+
+# True if we're running on FreeBSD.
+_IS_FREEBSD = sys.platform.startswith("freebsd")
+
+# True if we're running on MacOS.
+_IS_MACOS = sys.platform == "darwin"
+
+# True if we're in Github Actions.
+_IS_GITHUB_ACTIONS = os.environ.get("GITHUB_ACTIONS") is not None
+
+# True if we're using uvloop.
+_IS_UVLOOP = os.environ.get("SHELLOUS_LOOP_TYPE") == "uvloop"
+
+# The interactive prompt on PyPY3 is ">>>> ".
+if _IS_PYPY:
+    _PS1 = ">>>> "
+else:
+    _PS1 = ">>> "
+
+_NO_ECHO = cooked(echo=False)
+
+_requires_unix = pytest.mark.skipif(sys.platform == "win32", reason="requires unix")
+
+_requires_pty = pytest.mark.skipif(
+    _IS_UVLOOP or sys.platform == "win32",
+    reason="requires pty",
+)
+
+
+@_requires_pty
+async def test_prompt_python_pty():
+    "Test the prompt class with the Python REPL (PTY)."
+    cmd = sh(sys.executable).stdin(sh.CAPTURE).stdout(sh.CAPTURE).stderr(sh.STDOUT)
+
+    async with cmd.set(pty=_NO_ECHO) as run:
+        repl = Prompt(run, _PS1, default_timeout=3.0)
+
+        greeting = await repl.receive()
+        assert "Python" in greeting
+
+        result = await repl.send("print('abc')")
+        if _IS_MACOS and _IS_GITHUB_ACTIONS:
+            assert result == "print('abc')\nabc"
+        else:
+            assert result == "abc"
+
+        await repl.send("exit()")
+
+    assert run.result().exit_code == 0
+
+
+async def test_prompt_python_interactive():
+    "Test the prompt class with the Python REPL (non-PTY using -i)."
+    cmd = (
+        sh(sys.executable, "-i").stdin(sh.CAPTURE).stdout(sh.CAPTURE).stderr(sh.STDOUT)
+    )
+
+    async with cmd as run:
+        repl = Prompt(run, _PS1, default_timeout=3.0)
+
+        greeting = await repl.receive()
+        assert "Python" in greeting
+
+        result = await repl.send("print('abc')")
+        assert result == "abc"
+
+        await repl.send("exit()")
+
+    assert run.result().exit_code == 0
+
+
+@_requires_pty
+async def test_prompt_python_interactive_ps1():
+    "Test the Python REPL but change the prompt to something unique."
+    alt_ps1 = "????"
+    cmd = (
+        sh(sys.executable, "-i").stdin(sh.CAPTURE).stdout(sh.CAPTURE).stderr(sh.STDOUT)
+    )
+
+    async with cmd as run:
+        repl = Prompt(run, alt_ps1)
+
+        greeting = await repl.send(f"import sys; sys.ps1='{alt_ps1}'")
+        assert _PS1 in greeting
+
+        result = await repl.send("print('def')")
+        assert result == "def"
+
+        await repl.send("exit()")
+
+    assert run.result().exit_code == 0
+
+
+@_requires_pty
+async def test_prompt_python_timeout():
+    "Test the prompt class with the Python REPL."
+    cmd = (
+        sh(sys.executable, "-i").stdin(sh.CAPTURE).stdout(sh.CAPTURE).stderr(sh.STDOUT)
+    )
+
+    async with cmd as run:
+        repl = Prompt(run, _PS1)
+
+        greeting = await repl.receive()
+        assert "Python" in greeting
+
+        with pytest.raises(asyncio.TimeoutError):
+            await repl.send("import time; time.sleep(1)", timeout=0.1)
+
+        await repl.send("exit()")
+
+    assert run.result().exit_code == 0
+
+
+@_requires_pty
+async def test_prompt_python_missing_newline():
+    "Test the prompt class with the Python REPL."
+    cmd = (
+        sh(sys.executable, "-i").stdin(sh.CAPTURE).stdout(sh.CAPTURE).stderr(sh.STDOUT)
+    )
+
+    async with cmd as run:
+        repl = Prompt(run, _PS1)
+
+        greeting = await repl.receive()
+        assert "Python" in greeting
+
+        result = await repl.send("print(3, end='.')")
+        assert result == "3."
+
+        await repl.send("exit()")
+
+    assert run.result().exit_code == 0
+
+
+@_requires_pty
+async def test_prompt_unix_shell():
+    "Test the prompt class with a shell (PTY no echo)."
+    cmd = (
+        sh("sh")
+        .stdin(sh.CAPTURE)
+        .stdout(sh.CAPTURE)
+        .stderr(sh.STDOUT)
+        .set(pty=_NO_ECHO, inherit_env=False)
+    )
+
+    async with cmd.env(PS1="$", TERM="dumb") as run:
+        repl = Prompt(run, "$")
+
+        greeting = await repl.receive()
+        if _IS_FREEBSD:
+            # FIXME: FreeBSD is complaining that it can't access tty?
+            assert "job control" in greeting
+        else:
+            assert greeting == ""
+
+        result = await repl.send("echo 123")
+        if _IS_FREEBSD:
+            # FIXME: I don't understand why FreeBSD is still echoing?
+            assert result == "echo 123\n123"
+        else:
+            assert result == "123"
+
+        await repl.send("exit")
+
+    assert run.result().exit_code == 0
+
+
+@_requires_pty
+async def test_prompt_unix_shell_echo():
+    "Test the prompt class with a shell (PTY default cooked mode)."
+    cmd = (
+        sh("sh")
+        .stdin(sh.CAPTURE)
+        .stdout(sh.CAPTURE)
+        .stderr(sh.STDOUT)
+        .set(pty=True, inherit_env=False)
+    )
+
+    async with cmd.env(PS1="$", TERM="dumb") as run:
+        repl = Prompt(run, "$")
+
+        greeting = await repl.receive()
+        if _IS_FREEBSD:
+            # FIXME: FreeBSD is complaining that it can't access tty?
+            assert "job control" in greeting
+        else:
+            assert greeting == ""
+
+        result = await repl.send("echo 123")
+        if _IS_ALPINE:
+            # Alpine is including terminal escape chars.
+            assert result == "\x1b[6necho 123\n123"
+        else:
+            assert result == "echo 123\n123"
+
+        await repl.send("exit")
+
+    assert run.result().exit_code == 0
+
+
+@_requires_unix
+async def test_prompt_unix_shell_interactive():
+    "Test the prompt class with an interactive shell (non-PTY, forced)."
+    cmd = (
+        sh("sh", "-i")
+        .stdin(sh.CAPTURE)
+        .stdout(sh.CAPTURE)
+        .stderr(sh.STDOUT)
+        .set(inherit_env=False)
+    )
+
+    async with cmd.env(PS1="$", TERM="dumb") as run:
+        repl = Prompt(run, "$")
+
+        greeting = await repl.receive()
+        assert "job control" in greeting  # expect message about job control
+
+        result = await repl.send("echo 123")
+
+        if _IS_MACOS and sys.version_info[:2] >= (3, 10) and not _IS_UVLOOP:
+            # On MacOS with Python 3.10 or later, the result is 'echo 123\n123'.
+            # (Unless we are using uvloop.) Result is "123" on MacOS in
+            # Python 3.9, or when using uvloop with later Python versions.
+            assert result == "echo 123\n123"
+        else:
+            assert result == "123"
+
+        await repl.send("exit")
+
+    assert run.result().exit_code == 0
+
+
+async def test_prompt_asyncio_repl():
+    "Test the prompt class with the asyncio REPL."
+    cmd = (
+        sh(sys.executable, "-m", "asyncio")
+        .stdin(sh.CAPTURE)
+        .stdout(sh.CAPTURE)
+        .stderr(sh.STDOUT)
+    )
+
+    async with cmd as run:
+        repl = Prompt(run, ">>> ")
+
+        greeting = await repl.receive()
+        assert "asyncio" in greeting
+
+        extra = await repl.receive()
+        assert "import asyncio" in extra
+
+        result = await repl.send("print('hello')")
+        assert result == "hello"
+
+        repl.close()
+
+    assert run.result().exit_code == 0
diff --git a/tests/test_readme.py b/tests/test_readme.py
index 9378609f..1da670d0 100644
--- a/tests/test_readme.py
+++ b/tests/test_readme.py
@@ -8,91 +8,52 @@
 import pytest
 
 from shellous import sh
-from shellous.harvest import harvest_results
 from shellous.log import LOGGER
+from shellous.prompt import Prompt
 
-_PROMPT = ">>> "
-_PROMPT_ELLIPSIS = "... "
-
-
-class Prompt:
-    "Utility class to help with an interactive prompt."
-
-    def __init__(self, stdin, stdout, errbuf):
-        self.stdin = stdin
-        self.stdout = stdout
-        self.errbuf = errbuf
-        self.prompt_bytes = _PROMPT.encode("utf-8")
-
-    async def prompt(self, input_text=""):
-        "Write some input text to stdin, then await the response."
-        if input_text:
-            self.stdin.write(input_text.encode("utf-8") + b"\n")
-
-        # Drain our write to stdin, and wait for prompt from stdout.
-        cancelled, (out, _) = await harvest_results(
-            self.stdout.readuntil(self.prompt_bytes),
-            self.stdin.drain(),
-            timeout=5.0,
-        )
-        if cancelled:
-            raise asyncio.CancelledError()
-
-        # If there are ellipsis bytes in the beginning of out, remove them.
-        assert isinstance(out, bytes)
-        out = re.sub(rb"^(?:\.\.\. )+", b"", out)
-
-        # Combine stderr and stdout, then clear stderr.
-        buf = self.errbuf + out
-        self.errbuf.clear()
-
-        # Clean up the output to remove the prompt, then return as string.
-        buf = buf.replace(b"\r\n", b"\n")
-        assert buf.endswith(self.prompt_bytes)
-        promptlen = len(self.prompt_bytes)
-        buf = buf[0:-promptlen].rstrip(b"\n")
-
-        return buf.decode("utf-8")
+# Custom Python REPL prompt.
+_PROMPT = "<+=+=+>"
 
 
 async def run_asyncio_repl(cmds, logfile=None):
     "Helper function to run the asyncio REPL and feed it commands."
-    errbuf = bytearray()
     repl = (
         sh(sys.executable, "-m", "asyncio")
         .stdin(sh.CAPTURE)
         .stdout(sh.CAPTURE)
-        .stderr(errbuf)
-        .set(_return_result=True, inherit_env=False)
+        .stderr(sh.STDOUT)
+        .set(inherit_env=False)
         .env(**_current_env())
     )
 
     async with repl as run:
-        assert run.stdin is not None
+        prompt = Prompt(run, _PROMPT, default_timeout=5.0)
 
-        p = Prompt(run.stdin, run.stdout, errbuf)
-        await p.prompt()
+        # Customize the python REPL prompt to make it easier to detect. The
+        # initial output of the REPL will include the old ">>> " prompt which
+        # we can ignore.
+        await prompt.send(f"import sys; sys.ps1 = '{_PROMPT}'; sys.ps2 = ''")
 
         # Optionally redirect logging to a file.
-        await p.prompt("import shellous.log, logging")
+        await prompt.send("import shellous.log, logging")
         if logfile:
-            await p.prompt("shellous.log.LOGGER.setLevel(logging.DEBUG)")
-            await p.prompt(
+            await prompt.send("shellous.log.LOGGER.setLevel(logging.DEBUG)")
+            await prompt.send(
                 f"logging.basicConfig(filename='{logfile}', level=logging.DEBUG)"
             )
         else:
             # I don't want random logging messages to confuse the output.
-            await p.prompt("shellous.log.LOGGER.setLevel(logging.ERROR)")
+            await prompt.send("shellous.log.LOGGER.setLevel(logging.ERROR)")
 
         output = []
         for cmd in cmds:
             LOGGER.info("  repl: %r", cmd)
-            output.append(await p.prompt(cmd))
+            output.append(await prompt.send(cmd))
             # Give tasks a chance to get started.
             if ".create_task(" in cmd:
                 await asyncio.sleep(0.1)
 
-        run.stdin.close()
+        prompt.close()
 
     result = run.result()
     assert result.exit_code == 0
@@ -225,12 +186,12 @@ def _separate_cmds_and_outputs(lines):
     output = ""
 
     for line in lines:
-        if line.startswith(_PROMPT):
+        if line.startswith(">>> "):
             if cmd is not None:
                 yield (cmd, output.rstrip("\n"))
             cmd = line[4:]
             output = ""
-        elif line.startswith(_PROMPT_ELLIPSIS):
+        elif line.startswith("... "):
             assert cmd
             cmd += "\n" + line[4:]
         else: