improved exec and run monitoring

Author: pavelrevak

mpytool/mpy_comm.py

@@ -176,16 +176,18 @@ def reset_state(self):
         self._repl_mode = None
         self._raw_paste_supported = None
 
-    def exec(self, command, timeout=None):
+    def exec(self, command, timeout=None, stream=False):
         """Execute command
 
         Arguments:
             command: command to execute
             timeout: maximum waiting time for result (None = wait forever),
                 0 = submit only (send code, don't wait for output)
+            stream: False = return bytes, True = yield chunks,
+                file-like object = write chunks to it and return b''
 
         Returns:
-            command STDOUT result
+            bytes result, generator, or b'' (when stream is file-like)
 
         Raises:
             CmdError when command return error
@@ -199,6 +201,10 @@ def exec(self, command, timeout=None):
         if timeout == 0:
             self._repl_mode = False
             return b''
+        if stream is True:
+            return self._exec_stream_result(command, timeout)
+        if stream:
+            return self._exec_stream_to(command, timeout, stream)
         result = self._conn.read_until(CTRL_D, timeout)
         if result:
             self._log.info('RES: %s', bytes(result))
@@ -280,6 +286,66 @@ def _wait_for_paste_complete(self, timeout):
             if byte == CTRL_D:
                 break
 
+    def _read_stdout_chunks(self, timeout):
+        """Read stdout chunks until CTRL_D marker.
+
+        Yields (chunk, is_last) tuples. Timeout is total wall time.
+        """
+        start_time = _time.time()
+        while True:
+            if timeout is not None:
+                remaining = timeout - (_time.time() - start_time)
+                if remaining <= 0:
+                    raise _conn.Timeout("Execution timeout")
+            self._conn._read_to_buffer(wait_timeout=0.001)
+            idx = self._conn._buffer.find(CTRL_D)
+            if idx != -1:
+                chunk = bytes(self._conn._buffer[:idx])
+                del self._conn._buffer[:idx + len(CTRL_D)]
+                if chunk:
+                    yield chunk
+                return
+            if self._conn._buffer:
+                yield bytes(self._conn._buffer)
+                self._conn._buffer.clear()
+
+    def _check_stderr(self, command, result, timeout, start_time):
+        """Read stderr and raise CmdError if non-empty."""
+        if result:
+            self._log.info('RES: %s', bytes(result))
+        err_timeout = 5 if timeout is None else max(
+            1, timeout - (_time.time() - start_time))
+        err = self._conn.read_until(CTRL_D + b'>', err_timeout)
+        if err:
+            raise CmdError(command, bytes(result), err)
+
+    def _exec_stream_result(self, command, timeout):
+        """Read execution result as streaming generator."""
+        start_time = _time.time()
+        result = bytearray()
+        for chunk in self._read_stdout_chunks(timeout):
+            result.extend(chunk)
+            yield chunk
+        self._check_stderr(command, result, timeout, start_time)
+
+    def _exec_stream_to(self, command, timeout, out):
+        """Read execution result, writing chunks to file-like object.
+
+        Supports both binary and text streams (io.TextIOBase).
+        """
+        import io as _io
+        start_time = _time.time()
+        result = bytearray()
+        text_mode = isinstance(out, _io.TextIOBase)
+        for chunk in self._read_stdout_chunks(timeout):
+            result.extend(chunk)
+            out.write(
+                chunk.decode('utf-8', 'backslashreplace')
+                if text_mode else chunk)
+            out.flush()
+        self._check_stderr(command, result, timeout, start_time)
+        return b''
+
     def _read_execution_result(self, command, timeout):
         """Read and parse execution result."""
         result = self._conn.read_until(CTRL_D, timeout)
@@ -290,15 +356,17 @@ def _read_execution_result(self, command, timeout):
             raise CmdError(command.decode('utf-8', errors='replace'), result, err)
         return result
 
-    def exec_raw_paste(self, command, timeout=5):
+    def exec_raw_paste(self, command, timeout=5, stream=False):
         """Execute code via raw-paste mode (flow-controlled, less RAM).
 
         Arguments:
             command: code to execute (str or bytes)
             timeout: max wait time (0 = submit only, don't wait for output)
+            stream: False = return bytes, True = yield chunks,
+                file-like object = write chunks to it and return b''
 
         Returns:
-            command STDOUT result
+            bytes result, generator, or b'' (when stream is file-like)
 
         Raises:
             CmdError: command execution error
@@ -338,26 +406,32 @@ def exec_raw_paste(self, command, timeout=5):
             self._repl_mode = False
             return b''
 
+        if stream is True:
+            return self._exec_stream_result(command, timeout)
+        if stream:
+            return self._exec_stream_to(command, timeout, stream)
         return self._read_execution_result(command, timeout)
 
-    def try_raw_paste(self, command, timeout=5):
+    def try_raw_paste(self, command, timeout=5, stream=False):
         """Try raw-paste mode, fall back to regular exec if not supported.
 
         Arguments:
             command: command to execute
             timeout: maximum waiting time for result
+            stream: False = return bytes, True = yield chunks,
+                file-like object = write chunks to it and return b''
 
         Returns:
-            command STDOUT result
+            bytes result, generator, or b'' (when stream is file-like)
         """
         # If we know raw-paste is not supported, skip it
         if self._raw_paste_supported is False:
-            return self.exec(command, timeout)
+            return self.exec(command, timeout, stream=stream)
 
         try:
-            return self.exec_raw_paste(command, timeout)
+            return self.exec_raw_paste(command, timeout, stream=stream)
         except MpyError as e:
             if "not supported" in str(e):
                 self._log.info("Raw-paste not supported, using regular exec")
-                return self.exec(command, timeout)
+                return self.exec(command, timeout, stream=stream)
             raise

mpytool/mpytool.py

@@ -921,24 +921,39 @@ def _dispatch_repl(self, commands, is_last_group):
 
     @command('exec', 'Execute Python code on device.')
     @argument('code', help='Python code to execute')
+    @option('-t', '--timeout', type=float, default=None,
+        help='timeout in seconds (0 = fire-and-forget)')
     def _dispatch_exec(self, commands, is_last_group):
-        args = _make_parser(self._dispatch_exec).parse_args([commands.pop(0)])
+        args = _make_parser(self._dispatch_exec).parse_args(commands)
+        commands.clear()
         self.verbose(f"EXEC: {args.code}", 1)
-        result = self.mpy.comm.exec(args.code)
-        if result:
-            print(result.decode('utf-8', 'backslashreplace'), end='')
+        timeout = args.timeout
+        if timeout == 0:
+            self.mpy.comm.try_raw_paste(args.code, timeout=0)
+            return
+        out = getattr(_sys.stdout, 'buffer', _sys.stdout)
+        self.mpy.comm.try_raw_paste(
+            args.code, timeout=timeout, stream=out)
 
     @command('run', 'Run local Python file on device.')
     @argument('file', metavar='local_file', help='local .py file')
+    @option('-t', '--timeout', type=float, default=None,
+        help='timeout in seconds (0 = fire-and-forget)')
     def _dispatch_run(self, commands, is_last_group):
-        arg_list = [commands.pop(0)] if commands else []
-        args = _make_parser(self._dispatch_run).parse_args(arg_list)
+        args = _make_parser(self._dispatch_run).parse_args(commands)
+        commands.clear()
         if not _os.path.isfile(args.file):
             raise ParamsError(f"file not found: {args.file}")
         with open(args.file, 'rb') as f:
             code = f.read()
         self.verbose(f"RUN: {args.file} ({len(code)} bytes)", 1)
-        self.mpy.comm.try_raw_paste(code, timeout=0)
+        timeout = args.timeout
+        if timeout == 0:
+            self.mpy.comm.try_raw_paste(code, timeout=0)
+            return
+        out = getattr(_sys.stdout, 'buffer', _sys.stdout)
+        self.mpy.comm.try_raw_paste(
+            code, timeout=timeout, stream=out)
 
     def _get_editor(self, editor_arg=None):
         """Get editor from --editor, $VISUAL, or $EDITOR"""

tests/test_mpy.py

@@ -562,6 +562,149 @@ def test_try_raw_paste_timeout_zero_fallback_exec(self):
         self.assertFalse(self.comm._repl_mode)
 
 
+class TestExecStream(unittest.TestCase):
+    """Tests for exec() stream parameter (False, True, file-like object)"""
+
+    def setUp(self):
+        self.mock_conn = Mock()
+        self.mock_conn.read_until.return_value = b''
+        self.mock_conn.flush.return_value = None
+        self.comm = mpy_comm.MpyComm(self.mock_conn)
+        self.comm._repl_mode = True
+
+    def _setup_exec_response(self, stdout_data, stderr_data=b''):
+        """Configure mock_conn.read_until to return stdout then stderr."""
+        self.mock_conn.read_until.side_effect = [
+            b'',           # read_until(b'OK')
+            stdout_data,   # read_until(CTRL_D) - stdout
+            stderr_data,   # read_until(CTRL_D + b'>') - stderr
+        ]
+
+    def _setup_stream_buffer(self, data):
+        """Configure mock_conn._buffer for streaming reads.
+
+        Simulates data arriving in _buffer, terminated by CTRL_D + CTRL_D>.
+        """
+        buf = bytearray(data + mpy_comm.CTRL_D)
+        self.mock_conn._buffer = buf
+        self.mock_conn._read_to_buffer.return_value = False
+        # After stdout CTRL_D is consumed, read_until for stderr
+        self.mock_conn.read_until.side_effect = [
+            b'',   # read_until(b'OK')
+            b'',   # read_until(CTRL_D + b'>') - no error
+        ]
+
+    # --- stream=False (default) ---
+
+    def test_default_returns_bytes(self):
+        """exec(cmd) returns bytes result"""
+        self._setup_exec_response(b'hello\r\n')
+        result = self.comm.exec("print('hello')")
+        self.assertEqual(result, b'hello\r\n')
+
+    def test_default_returns_empty_bytes(self):
+        """exec(cmd) with no output returns empty bytes"""
+        self._setup_exec_response(b'')
+        result = self.comm.exec("x=1")
+        self.assertEqual(result, b'')
+
+    # --- stream=True (generator) ---
+
+    def test_yield_returns_generator(self):
+        """exec(cmd, stream=True) returns a generator"""
+        import types
+        self._setup_stream_buffer(b'data')
+        gen = self.comm.exec("print('data')", stream=True)
+        self.assertIsInstance(gen, types.GeneratorType)
+        # Consume to avoid ResourceWarning
+        list(gen)
+
+    def test_yield_produces_data(self):
+        """exec(cmd, stream=True) yields all data"""
+        self._setup_stream_buffer(b'hello world')
+        chunks = list(self.comm.exec("print('hello world')", stream=True))
+        self.assertEqual(b''.join(chunks), b'hello world')
+
+    def test_yield_empty_output(self):
+        """exec(cmd, stream=True) with no output yields nothing"""
+        self._setup_stream_buffer(b'')
+        chunks = list(self.comm.exec("x=1", stream=True))
+        self.assertEqual(chunks, [])
+
+    def test_yield_raises_cmd_error(self):
+        """exec(cmd, stream=True) raises CmdError on stderr"""
+        buf = bytearray(b'partial' + mpy_comm.CTRL_D)
+        self.mock_conn._buffer = buf
+        self.mock_conn._read_to_buffer.return_value = False
+        self.mock_conn.read_until.side_effect = [
+            b'',                    # read_until(b'OK')
+            b'NameError: xxx',      # read_until(CTRL_D + b'>') - error
+        ]
+        with self.assertRaises(mpy_comm.CmdError):
+            list(self.comm.exec("xxx", stream=True))
+
+    # --- stream=binary file object ---
+
+    def test_binary_stream_writes_bytes(self):
+        """exec(cmd, stream=BytesIO) writes bytes to stream"""
+        import io
+        self._setup_stream_buffer(b'binary output')
+        out = io.BytesIO()
+        result = self.comm.exec("print('binary output')", stream=out)
+        self.assertEqual(result, b'')
+        self.assertEqual(out.getvalue(), b'binary output')
+
+    def test_binary_stream_empty_output(self):
+        """exec(cmd, stream=BytesIO) with no output writes nothing"""
+        import io
+        self._setup_stream_buffer(b'')
+        out = io.BytesIO()
+        self.comm.exec("x=1", stream=out)
+        self.assertEqual(out.getvalue(), b'')
+
+    # --- stream=text file object ---
+
+    def test_text_stream_writes_str(self):
+        """exec(cmd, stream=StringIO) writes decoded str to stream"""
+        import io
+        self._setup_stream_buffer(b'text output')
+        out = io.StringIO()
+        result = self.comm.exec("print('text output')", stream=out)
+        self.assertEqual(result, b'')
+        self.assertEqual(out.getvalue(), 'text output')
+
+    def test_text_stream_decodes_utf8(self):
+        """exec(cmd, stream=StringIO) decodes UTF-8 correctly"""
+        import io
+        self._setup_stream_buffer('príliš žluťoučký'.encode('utf-8'))
+        out = io.StringIO()
+        self.comm.exec("print('...')", stream=out)
+        self.assertEqual(out.getvalue(), 'príliš žluťoučký')
+
+    def test_text_stream_handles_invalid_utf8(self):
+        """exec(cmd, stream=StringIO) uses backslashreplace for invalid bytes"""
+        import io
+        self._setup_stream_buffer(b'ok\xff\xfebad')
+        out = io.StringIO()
+        self.comm.exec("cmd", stream=out)
+        self.assertIn('ok', out.getvalue())
+        self.assertIn('bad', out.getvalue())
+
+    def test_text_stream_raises_cmd_error(self):
+        """exec(cmd, stream=StringIO) raises CmdError on stderr"""
+        import io
+        buf = bytearray(b'out' + mpy_comm.CTRL_D)
+        self.mock_conn._buffer = buf
+        self.mock_conn._read_to_buffer.return_value = False
+        self.mock_conn.read_until.side_effect = [
+            b'',
+            b'SyntaxError',
+        ]
+        out = io.StringIO()
+        with self.assertRaises(mpy_comm.CmdError):
+            self.comm.exec("bad", stream=out)
+
+
 class TestCwdAndPathTracking(unittest.TestCase):
     """Tests for CWD and sys.path tracking for soft reset restore"""
 

tests/test_mpytool.py

@@ -712,25 +712,33 @@ def test_run_file_not_found(self):
         self.assertIn('not found', str(ctx.exception).lower())
 
     def test_run_sends_file_content(self):
-        """'run script.py' -> reads file and sends via try_raw_paste"""
+        """'run script.py' -> reads file and sends via try_raw_paste (streaming)"""
         self.tool.process_commands(['run', self.test_file])
-        self.tool._mpy.comm.try_raw_paste.assert_called_once_with(
-            b"print('hello')\n", timeout=0)
+        args, kwargs = self.tool._mpy.comm.try_raw_paste.call_args
+        self.assertEqual(args[0], b"print('hello')\n")
+        self.assertIsNone(kwargs['timeout'])
+        self.assertNotIn(kwargs['stream'], (False, True))  # file-like object
 
     def test_run_reads_binary_mode(self):
         """run reads file as bytes (rb), preserving encoding"""
         utf8_file = os.path.join(self.temp_dir, "utf8.py")
         with open(utf8_file, 'wb') as f:
             f.write("print('súbor')\n".encode('utf-8'))
         self.tool.process_commands(['run', utf8_file])
-        self.tool._mpy.comm.try_raw_paste.assert_called_once_with(
-            "print('súbor')\n".encode('utf-8'), timeout=0)
+        args, kwargs = self.tool._mpy.comm.try_raw_paste.call_args
+        self.assertEqual(args[0], "print('súbor')\n".encode('utf-8'))
 
     def test_run_with_command_separator(self):
         """'run script.py -- ls' -> run then ls"""
         self.tool.process_commands(['run', self.test_file])
         self.tool._mpy.comm.try_raw_paste.assert_called_once()
 
+    def test_run_fire_and_forget(self):
+        """'run script.py -t 0' -> fire-and-forget (timeout=0)"""
+        self.tool.process_commands(['run', self.test_file, '-t', '0'])
+        self.tool._mpy.comm.try_raw_paste.assert_called_once_with(
+            b"print('hello')\n", timeout=0)
+
 
 class TestPathCommand(unittest.TestCase):
     """Tests for path command dispatch"""