Refactor Subprocess, document expects

README.rst

@@ -7,8 +7,7 @@ Why?
 In Python we now have async subprocesses which allows to execute several
 subprocesses at the same time. The purpose of this library is to:
 
-- provide an acceptable asyncio subprocess wrapper for my syntaxic taste,
+- stream stderr and stdout in real time while capturing it,
-- can stream stderr and stdout in real time while capturing it,
 - real time output must be prefixed for when you execute several commands at
   the time so that you know which line is for which process, like with
   docker-compose logs,
@@ -16,6 +15,13 @@ subprocesses at the same time. The purpose of this library is to:
 
 This code was copy/pasted between projects and finally extracted on its own.
 
+Demo
+====
+
+.. image:: https://yourlabs.io/oss/shlax/-/raw/master/demo.png
+
+You will find the demo script in demo.py in this repository.
+
 Usage
 =====
 
@@ -30,6 +36,8 @@ Basic example, this will both stream output and capture it:
     proc = await Subprocess('echo hi').wait()
     print(proc.rc, proc.out, proc.err, proc.out_raw, proc.err_raw)
 
+Arguments may be a string command or a list of arguments.
+
 Longer
 ------
 
@@ -38,7 +46,7 @@ any of ``start()`` and ``wait()``, or both, explicitely:
 
 .. code-block:: python
 
-    proc = Subprocess('echo hi')
+    proc = Subprocess('echo', 'hi')
     await proc.start()  # start the process
     await proc.wait()   # wait for completion
 
@@ -87,17 +95,36 @@ will be applied line by line:
     }
     await asyncio.gather(*[
         Subprocess(
-            f'find {path}',
+            'find',
+            path,
             regexps=regexps,
+            shell=True,
         ).wait()
         for path in sys.path
     ])
 
+Automating input
+----------------
+
+You can pass a list of tuples of two bytestrings ``(regexp, characters_to_send)``:
+
+.. code-block:: python
+
+    proc = Proc(
+        'sh',
+        '-euc',
+        'echo "x?"; read x; echo x=$x; echo "z?"; read z; echo z=$z',
+        expects=[
+            (b'x?', b'y\n'),
+            (b'z?', b'w\n'),
+        ],
+    )
+    await proc.wait()
+    assert proc.out == 'x?\nx=y\nz?\nz=w'
+
 Where is the rest?
 ==================
 
-Shlax used to be the name of a much more ambitious poc-project that has been
+Shlax used to be the name of a much more ambitious poc-project, that you can
-extracted in two projects with clear boundaries, namely `sysplan
+still find in the ``OLD`` branch of this repository. Parts of it have been
-<https://yourlabs.io/oss/sysplan>`_ and `podplan
+extracted into smaller repositories.
-<https://yourlabs.io/oss/podplan>`_ which are still in alpha state, but Shlax
-as it is feature complete and stable.

demo.py

@@ -0,0 +1,26 @@
+import asyncio
+
+from shlax import Subprocess
+
+
+async def main():
+    colors = {
+        '^(.*).txt$': '{green}\\1.txt',
+        '^(.*).py$': '{bred}\\1.py',
+    }
+
+    await asyncio.gather(
+         Subprocess(
+            'sh', '-euc',
+            'for i in $(find .. | head); do echo $i; sleep .2; done',
+            regexps=colors,
+            prefix='parent',
+        ).wait(),
+        Subprocess(
+            'sh -euc "for i in $(find . | head); do echo $i; sleep .3; done"',
+            regexps=colors,
+            prefix='cwd',
+        ).wait()
+    )
+
+asyncio.run(main())

shlax/subprocess.py

@@ -1,24 +1,38 @@
 import asyncio
 import functools
+import os
 import re
+import shlex
 import sys
 
 from .colors import colors
 
 
-class SubprocessProtocol(asyncio.SubprocessProtocol):
+class SubprocessProtocol(asyncio.subprocess.SubprocessStreamProtocol):
-    def __init__(self, proc):
+    def __init__(self, proc, *args, **kwargs):
         self.proc = proc
-        self.output = bytearray()
+        super().__init__(*args, **kwargs)
+
+    def receive(self, data, raw, target):
+        raw.extend(data)
+        if not self.proc.quiet:
+            for line in self.proc.lines(data):
+                target.buffer.write(line)
+            target.flush()
 
     def pipe_data_received(self, fd, data):
         if fd == 1:
-            self.proc.stdout(data)
+            self.receive(data, self.proc.out_raw, self.proc.stdout)
         elif fd == 2:
-            self.proc.stderr(data)
+            self.receive(data, self.proc.err_raw, self.proc.stderr)
 
-    def process_exited(self):
+        if self.proc.expect_index < len(self.proc.expects):
-        self.proc.exit_future.set_result(True)
+            expected = self.proc.expects[self.proc.expect_index]
+            if re.match(expected[0], data):
+                self.stdin.write(expected[1])
+                event_loop = asyncio.get_event_loop()
+                asyncio.create_task(self.stdin.drain())
+                self.proc.expect_index += 1
 
 
 class Subprocess:
@@ -48,18 +62,19 @@ class Subprocess:
         quiet=None,
         prefix=None,
         regexps=None,
+        expects=None,
         write=None,
         flush=None,
+        stdout=None,
+        stderr=None,
     ):
-        if len(args) == 1 and ' ' in args[0]:
-            args = ['sh', '-euc', args[0]]
-
-        self.cmd = ' '.join(args)
         self.args = args
         self.quiet = quiet if quiet is not None else False
         self.prefix = prefix
-        self.write = write or sys.stdout.buffer.write
+        self.stdout = stdout or sys.stdout
-        self.flush = flush or sys.stdout.flush
+        self.stderr = stderr or sys.stderr
+        self.expects = expects or []
+        self.expect_index = 0
         self.started = False
         self.waited = False
         self.out_raw = bytearray()
@@ -75,19 +90,41 @@ class Subprocess:
                 self.regexps[search] = replace
 
     async def start(self, wait=True):
-        # Get a reference to the event loop as we plan to use
+        if len(self.args) == 1 and not os.path.exists(self.args[0]):
-        # low-level APIs.
+            args = shlex.split(self.args[0])
+        else:
+            args = self.args
+
+        if not self.quiet:
+            message = b''.join([
+                self.colors.bgray.encode(),
+                b'+ ',
+                shlex.join(args).replace('\n', '\\n').encode(),
+                self.colors.reset.encode(),
+            ])
+            for line in self.lines(message, highlight=False):
+                self.stdout.buffer.write(line)
+            self.stdout.flush()
+
+        # The following is a copy of what asyncio.subprocess_exec and
+        # asyncio.create_subprocess_exec do except we inject our own
+        # SubprocessStreamProtocol subclass: it might need an update as new
+        # python releases come out.
         loop = asyncio.get_running_loop()
 
-        self.exit_future = asyncio.Future(loop=loop)
-
-        # Create the subprocess controlled by DateProtocol;
-        # redirect the standard output into a pipe.
         self.transport, self.protocol = await loop.subprocess_exec(
-            lambda: SubprocessProtocol(self),
+            lambda: SubprocessProtocol(
-            *self.args,
+                self,
-            stdin=None,
+                limit=asyncio.subprocess.streams._DEFAULT_LIMIT,
+                loop=loop,
+            ),
+            *args,
+            stdin=asyncio.subprocess.PIPE if self.expects else sys.stdin,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
         )
+
+        self.proc = asyncio.subprocess.Process(self.transport, self.protocol, loop)
         self.started = True
 
     async def wait(self, *args, **kwargs):
@@ -95,50 +132,35 @@ class Subprocess:
             await self.start()
 
         if not self.waited:
-            # Wait for the subprocess exit using the process_exited()
+            await self.proc.communicate()
-            # method of the protocol.
+            self.rc = self.transport.get_returncode()
-            await self.exit_future
-
-            # Close the stdout pipe.
-            self.transport.close()
-
             self.waited = True
 
         return self
 
-    def stdout(self, data):
+    @property
-        self.out_raw.extend(data)
-        if not self.quiet:
-            self.output(data)
-
-    def stderr(self, data):
-        self.err_raw.extend(data)
-        if not self.quiet:
-            self.output(data)
-
-    @functools.cached_property
     def out(self):
+        if self.waited:
+            if '_out_cached' not in self.__dict__:
+                self._out_cached = self.out_raw.decode().strip()
+            return self._out_cached
         return self.out_raw.decode().strip()
 
-    @functools.cached_property
+    @property
     def err(self):
+        if self.waited:
+            if '_err_cached' not in self.__dict__:
+                self._err_cached = self.err_raw.decode().strip()
+            return self._err_cached
         return self.err_raw.decode().strip()
 
-    @functools.cached_property
+    def lines(self, data, highlight=True):
-    def rc(self):
-        return self.transport.get_returncode()
-
-    def output(self, data, highlight=True, flush=True):
         for line in data.strip().split(b'\n'):
             line = [self.highlight(line) if highlight else line]
             if self.prefix:
                 line = self.prefix_line() + line
             line.append(b'\n')
-            line = b''.join(line)
+            yield b''.join(line)
-            self.write(line)
-
-        if flush:
-            self.flush()
 
     def highlight(self, line, highlight=True):
         if not highlight or (

tests/test_proc.py

@@ -8,9 +8,9 @@ from shlax import Proc
 @pytest.mark.parametrize(
     'args',
     (
-        ['sh', '-c', 'echo hi'],
         ['echo hi'],
         ['sh -c "echo hi"'],
+        ['sh', '-c', 'echo hi'],
     )
 )
 async def test_proc(args):
@@ -36,13 +36,11 @@ async def test_wait_unbound():
 @pytest.mark.asyncio
 async def test_rc_1():
     proc = await Proc(
-        'NON EXISTING COMMAND',
+        'sh', '-euc', 'NON EXISTING COMMAND',
-        write=Mock(),
+        quiet=True,
     ).wait()
     assert proc.rc != 0
-    proc.write.assert_called_once_with(
+    assert proc.err == 'sh: line 1: NON: command not found'
-        b'sh: line 1: NON: command not found\x1b[0m\n'
-    )
 
 
 @pytest.mark.asyncio
@@ -52,24 +50,34 @@ async def test_prefix():
     """
     Proc.prefix_length = 0  # reset
 
-    write = Mock()
+    stdout = Mock()
     await Proc(
         'echo hi',
-        write=write,
+        stdout=stdout,
         prefix='test_prefix',
     ).wait()
     await Proc(
         'echo hi',
-        write=write,
+        stdout=stdout,
         prefix='test_prefix_1'
     ).wait()
     await Proc(
         'echo hi',
-        write=write,
+        stdout=stdout,
         prefix='test_prefix',
     ).wait()
 
-    assert write.mock_calls == [
+    assert stdout.buffer.write.mock_calls == [
+        call(
+            Proc.prefix_colors[0].encode()
+            + b'test_prefix '
+            + Proc.colors.reset.encode()
+            + b'| '
+            + Proc.colors.bgray.encode()
+            + b'+ echo hi'
+            + Proc.colors.reset.encode()
+            + b'\n'
+        ),
         call(
             Proc.prefix_colors[0].encode()
             + b'test_prefix '
@@ -78,6 +86,16 @@ async def test_prefix():
             + Proc.colors.reset.encode()
             + b'\n'
         ),
+        call(
+            Proc.prefix_colors[1].encode()
+            + b'test_prefix_1 '
+            + Proc.colors.reset.encode()
+            + b'| '
+            + Proc.colors.bgray.encode()
+            + b'+ echo hi'
+            + Proc.colors.reset.encode()
+            + b'\n'
+        ),
         call(
             Proc.prefix_colors[1].encode()
             # padding has been added because of output1
@@ -87,6 +105,17 @@ async def test_prefix():
             + Proc.colors.reset.encode()
             + b'\n'
         ),
+        call(
+            Proc.prefix_colors[0].encode()
+            # padding has been added because of output1
+            + b'  test_prefix '
+            + Proc.colors.reset.encode()
+            + b'| '
+            + Proc.colors.bgray.encode()
+            + b'+ echo hi'
+            + Proc.colors.reset.encode()
+            + b'\n'
+        ),
         call(
             Proc.prefix_colors[0].encode()
             # padding has been added because of output1
@@ -104,10 +133,20 @@ async def test_prefix_multiline():
     Proc.prefix_length = 0  # reset
     proc = await Proc(
         'echo -e "a\nb"',
-        write=Mock(),
+        stdout=Mock(),
         prefix='test_prefix',
     ).wait()
-    assert proc.write.mock_calls == [
+    assert proc.stdout.buffer.write.mock_calls == [
+        call(
+            Proc.prefix_colors[0].encode()
+            + b'test_prefix '
+            + Proc.colors.reset.encode()
+            + b'| '
+            + Proc.colors.bgray.encode()
+            + b"+ echo -e 'a\\nb'"
+            + Proc.colors.reset.encode()
+            + b'\n'
+        ),
         call(
             Proc.prefix_colors[0].encode()
             + b'test_prefix '
@@ -135,12 +174,12 @@ async def test_highlight():
     """
     proc = await Proc(
         'echo hi',
-        write=Mock(),
+        stdout=Mock(),
         regexps={
             r'h([\w\d-]+)': 'h{cyan}\\1',
         }
     ).wait()
-    proc.write.assert_called_once_with(b'h\x1b[38;5;51mi\x1b[0m\n')
+    proc.stdout.buffer.write.assert_called_with(b'h\x1b[38;5;51mi\x1b[0m\n')
 
 
 @pytest.mark.asyncio
@@ -150,9 +189,40 @@ async def test_highlight_if_not_colored():
     """
     proc = await Proc(
         'echo -e h"\\e[31m"i',
-        write=Mock(),
+        stdout=Mock(),
         regexps={
             r'h([\w\d-]+)': 'h{cyan}\\1',
         }
     ).wait()
-    proc.write.assert_called_once_with(b'h\x1b[31mi\n')
+    proc.stdout.buffer.write.assert_called_with(b'h\x1b[31mi\n')
+
+
+@pytest.mark.asyncio
+async def test_expect():
+    proc = Proc(
+        'sh', '-euc',
+        'echo "x?"; read x; echo x=$x; echo "z?"; read z; echo z=$z',
+        expects=[
+            (b'x?', b'y\n'),
+            (b'z?', b'w\n'),
+        ],
+        quiet=True,
+    )
+    await proc.wait()
+    assert proc.out == 'x?\nx=y\nz?\nz=w'
+
+@pytest.mark.asyncio
+async def test_stderr():
+    proc = await Proc(
+        'sh',
+        '-euc',
+        'echo hi >&2',
+        stdout=Mock(),
+        stderr=Mock()
+    ).wait()
+
+    assert proc.err_raw == bytearray(b'hi\n')
+    assert proc.err == 'hi'
+    proc.stderr.buffer.write.assert_called_once_with(
+        f'hi{proc.colors.reset}\n'.encode()
+    )