lots of improvements; handling of continuation prompts; test with python repl

Author: jhidding

README.md

@@ -42,19 +42,19 @@ I like to work with [Chez Scheme](https://cisco.github.io/ChezScheme/). Suppose
 ```yaml
 #| file: test/scheme.yml
 config:
-    command: "scheme --eedisable"
-    first_prompt: "> "
-    change_prompt: "(waiter-prompt-string \"{key}>\")"
-    next_prompt: "{key}> "
-    strip_command: true
+  command: "scheme --eedisable"
+  first_prompt: "> "
+  change_prompt: '(waiter-prompt-string "{key}>")'
+  prompt: "{key}> "
+  continuation_prompt: ""
 commands:
-   - command: (* 6 7)
-   - command: |
+  - command: (* 6 7)
+  - command: |
       (define (fac n init)
         (if (zero? n)
           init
           (fac (- n 1) (* init n)))))
-   - command: (fac 10 1)
+  - command: (fac 10 1)
 ```
 
 Passing this to `repl-session`, it will start the Chez Scheme interpreter, waiting for the `>` prompt to appear. It then changes the prompt to a generated `uuid4` code, for instance `27e87a8a-742c-4501-b05d-b05814f5a010> `. This will make sure that we can't accidentally match something else for an interactive prompt (imagine we're generating some XML!). Since commands are also echoed to standard out, we need to strip them from the resulting output. Running this should give:
@@ -76,15 +76,27 @@ This looks very similar to the previous example:
 ```yaml
 #| file: test/lua.yml
 config:
-    command: "lua"
-    first_prompt: "> "
-    change_prompt: "_PROMPT = \"{key}> \""
-    next_prompt: "{key}> "
-    strip_command: true
-    strip_ansi: true
+  command: "lua"
+  first_prompt: "> "
+  change_prompt: '_PROMPT = "{key}> "; _PROMPT2 = "{key}+ "'
+  prompt: "{key}> "
+  continuation_prompt: "{key}\\+ "
+  strip_ansi: true
 commands:
-   - command: 6 * 7
-   - command: "\"Hello\" .. \", \" .. \"World!\""
+  - command: 6 * 7
+  - command: '"Hello" .. ", " .. "World!"'
+  - command: |
+      function fac(n, m)
+          if m == nil then
+              return fac(n, 1)
+          end
+          if n == 0 then
+              return m
+          else
+              return fac(n-1, m*n)
+          end
+      end
+  - command: fac(10)
 ```
 
 The Lua REPL is not so nice. It sends ANSI escape codes and those need to be filtered out.
@@ -98,6 +110,31 @@ repl-session < test/lua.yml | jq '.commands.[].output'
 "Hello, World!"
 ```
 
+### Python
+
+The Python REPL got a revision in version 3.13, with lots of colour and ANSI codes.
+
+```yaml
+#| file: test/python.yml
+config:
+  command: python -q
+  first_prompt: ">>>"
+  change_prompt: 'import sys; sys.ps1 = "{key}>>> "; sys.ps2 = "{key}+++ "'
+  prompt: "{key}>>> "
+  continuation_prompt: "{key}\\+\\+\\+ "
+  environment:
+    NO_COLOR: "1"
+commands:
+  - command: print("Hello, World!")
+  - command: 6 * 7
+  - command: |
+      def fac(n):
+          for i in range(1, n):
+              n *= i
+          return n
+  - command: fac(10)
+```
+
 ## Input/Output structure
 
 The user can configure how the REPL is called and interpreted.
@@ -119,14 +156,17 @@ class ReplConfig(msgspec.Struct):
             output; useful if the REPL echoes your input before answering.
         timeout (float): Command timeout for this session in seconds.
     """
+
     command: str
     first_prompt: str
     change_prompt: str
-    next_prompt: str
-    append_newline: bool = True
-    strip_command: bool = False
+    prompt: str
+    continuation_prompt: str | None = None
     strip_ansi: bool = False
+    environment: dict[str, str] = msgspec.field(default_factory=dict)
     timeout: float = 5.0
+
+
 ```
 
 Then, a session is a list of commands. Each command should be a UTF-8 string, and we allow to attach some meta-data like expected MIME type for the output. We can also pass an expected output in the case of a documentation test. If `output` was already given on the input, it is  moved to `expected`. This way it becomes really easy to setup regression tests on your documentation. Just rerun on the generated output file.
@@ -142,6 +182,7 @@ class ReplCommand(msgspec.Struct):
         output (str | None): evaluated output.
         expected (str | None): expected output.
     """
+
     command: str
     output_type: str = "text/plain"
     output: str | None = None
@@ -155,8 +196,11 @@ class ReplSession(msgspec.Struct):
         config (ReplConfig): Config for setting up a REPL session.
         commands (list[ReplCommand]): List of commands in the session.
     """
+
     config: ReplConfig
     commands: list[ReplCommand]
+
+
 ```
 
 ## License and contribution

docs/index.md

@@ -6,39 +6,85 @@ The core of the implementation is handled by the `pexpect` library. We have a sm
 
 ```python
 #| id: repl-contextmanager
+def spawn(config: ReplConfig):
+    child: pexpect.spawn[str] = pexpect.spawn(
+        config.command,
+        timeout=config.timeout,
+        echo=False,
+        encoding="utf-8",
+        env=config.environment,
+    )
+    return child
+
 
 @contextmanager
-def repl(config: ReplConfig) -> Generator[Callable[[str], str]]:
+def repl(config: ReplConfig) -> Generator[Callable[[str], str | None]]:
     key = str(uuid.uuid4())
+    change_prompt_cmd = config.change_prompt.format(key=key)
+    prompt = config.prompt.format(key=key)
+    continuation_prompt = (
+        config.continuation_prompt.format(key=key)
+        if config.continuation_prompt is not None
+        else None
+    )
+
+    child: pexpect.spawn[str]
+    with spawn(config) as child:
+        _ = child.expect(config.first_prompt)
+        _ = child.sendline(change_prompt_cmd)
+        # if config.strip_command:
+        #    child.expect(key)
+        _ = child.expect(prompt)
+
+        def send(msg: str) -> str | None:
+            nonlocal prompt, continuation_prompt, change_prompt_cmd
+            lines = msg.splitlines()
+            answer: list[str] = []
+
+            still_waiting: bool = True
+            for line in lines:
+                logging.debug("sending: %s", line)
+                _ = child.sendline(line)
+                if continuation_prompt is not None:
+                    logging.debug("waiting for prompt or continuation")
+                    _ = child.expect(
+                        f"(?P<cont>{continuation_prompt})|(?P<norm>{prompt})"
+                    )
+                    if not isinstance(child.match, re.Match):
+                        continue
+                    if child.match.group("cont") is not None:
+                        logging.debug("continuation")
+                        still_waiting = True
+                    else:
+                        logging.debug("done: %s", child.before)
+                        if child.before is not None:
+                            answer.append(child.before)
+                        still_waiting = False
+                else:
+                    logging.debug("waiting for prompt")
+                    _ = child.expect(prompt)
+                    if child.before:
+                        answer.append(child.before)
+                    still_waiting = False
+
+            if still_waiting:
+                _ = child.sendline("")
+                _ = child.expect(prompt)
+                if child.before:
+                    answer.append(child.before)
+
+            if not answer:
+                return None
 
-    with pexpect.spawn(config.command, timeout=config.timeout) as child:
-        child.expect(config.first_prompt)
-        change_prompt_cmd = config.change_prompt.format(key=key)
-        if config.append_newline:
-            change_prompt_cmd = change_prompt_cmd + "\n"
-        child.send(change_prompt_cmd)
-        if config.strip_command:
-            child.expect(key)
-        prompt = config.next_prompt.format(key=key)
-        child.expect(prompt)
-
-        def send(msg: str) -> str:
-            if config.append_newline:
-                msg = msg + "\n"
-            child.send(msg)
-            child.expect("(.*)" + prompt)
-
-            answer = child.match[1].decode()
             if config.strip_ansi:
-                ansi_escape = re.compile(r'(\u001b\[|\x1B\[)[0-?]*[ -\/]*[@-~]')
-                answer = ansi_escape.sub("", answer)
-            if config.strip_command:
-                answer = answer.strip().replace("\r", "")
-                return answer.removeprefix(msg)
-            else:
-                return child.match[1].decode()
+                ansi_escape = re.compile(r"(\u001b\[|\x1B\[)[0-?]*[ -\/]*[@-~]")
+                return ansi_escape.sub("", answer[-1].strip())
+
+            return answer[-1].strip()
 
         yield send
+
+
 ```
 
 We use this to run a session. The session is modified in place.
@@ -54,6 +100,8 @@ def run_session(session: ReplSession):
             cmd.expected = expected
 
     return session
+
+
 ```
 
 ### I/O
@@ -62,29 +110,32 @@ I/O is handled by `msgspec`.
 
 ```python
 #| id: io
-
 def read_session(port: IO[str] = sys.stdin) -> ReplSession:
     data: str = port.read()
     return msgspec.yaml.decode(data, type=ReplSession)
 
 
 def write_session(session: ReplSession, port: IO[str] = sys.stdout):
     data = msgspec.json.encode(session)
-    port.write(data.decode())
+    _ = port.write(data.decode())
+
+
 ```
 
 ##  Imports
 
 ```python
 #| id: imports
 # from datetime import datetime, tzinfo
-from typing import IO
+from typing import IO, cast
 from collections.abc import Generator, Callable
+
 # import re
 from contextlib import contextmanager
 import uuid
 import sys
 import re
+import logging
 
 import pexpect
 import msgspec
@@ -103,17 +154,21 @@ __version__ = importlib.metadata.version("repl-session")
 `repl-session` is a command-line tool to evaluate a given session
 in any REPL, and store the results.
 """
+
 <<imports>>
 
+
 <<input-data>>
 
+
 <<repl-contextmanager>>
 <<run-session>>
 <<io>>
 
 
 @argh.arg("-v", "--version", help="show version and exit")
-def repl_session(version: bool = False):
+@argh.arg("-l", "--log-enable", help="show debugging output")
+def repl_session(version: bool = False, log_enable: bool = False):
     """
     repl-session runs a REPL session, reading JSON from standard input and
     writing to standard output. Both the input and output follow the same
@@ -123,9 +178,14 @@ def repl_session(version: bool = False):
         print(f"repl-session {__version__}")
         sys.exit(0)
 
+    if log_enable:
+        logging.basicConfig(level=logging.DEBUG)
+
     write_session(run_session(read_session()))
 
 
 def main():
     argh.dispatch_command(repl_session)
+
+
 ```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "repl-session"
-version = "0.1.0"
+version = "0.2.0"
 description = "Runs a session with a Read-Eval-Print-Loop (REPL). Takes YAML as input gives JSON as output."
 readme = "README.md"
 classifiers = [
@@ -24,6 +24,8 @@ dev = [
     "entangled-cli>=2.1.13",
     "mkdocs-macros-plugin>=1.3.9",
     "mkdocs-material>=9.6.20",
+    "python-lsp-server>=1.13.1",
+    "ruff>=0.13.2",
     "ty>=0.0.1a21",
 ]