Skip to content
Back to top

# Source code for mindroot.lib.streamcmd

    import asyncio
    import sys
    from typing import Optional, Callable, Any, List, Dict, AsyncGenerator
    COLORS = {'reset': '\x1b[0m', 'green': '\x1b[32m', 'yellow': '\x1b[33m', 'blue': '\x1b[34m', 'red': '\x1b[31m'}

    async def read_stream(stream: asyncio.StreamReader, callback: Callable[[str], Any]):
        """Read from stream line by line and call the callback for each line."""
        while True:
            line = await stream.readline()
            if not line:
                break
            callback(line.decode('utf-8', errors='replace'))

    async def run_command_with_streaming(cmd: List[str], stdout_callback: Callable[[str], Any], stderr_callback: Callable[[str], Any], cwd: Optional[str]=None, env: Optional[dict]=None) -> int:
        """Run a command asynchronously and stream its output.

        Args:
            cmd: Command to run as a list of strings
            stdout_callback: Callback for stdout lines
            stderr_callback: Callback for stderr lines
            cwd: Working directory for the command
            env: Environment variables for the command

        Returns:
            Exit code of the command
        """
        try:
            process = await asyncio.create_subprocess_exec(*cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=cwd, env=env)
            stdout_task = asyncio.create_task(read_stream(process.stdout, stdout_callback))
            stderr_task = asyncio.create_task(read_stream(process.stderr, stderr_callback))
            await asyncio.gather(stdout_task, stderr_task)
            exit_code = await process.wait()
            return exit_code
        except Exception as e:
            return 1

    async def stream_command_as_events(cmd: List[str], cwd: Optional[str]=None, env: Optional[dict]=None) -> AsyncGenerator[Dict[str, str], None]:
        """Run a command and yield its output as events.

        Args:
            cmd: Command to run as a list of strings
            cwd: Working directory for the command
            env: Environment variables for the command

        Yields:
            Events with type and data
        """
        yield {'event': 'message', 'data': f"Running command: {' '.join(cmd)}"}
        output_queue = asyncio.Queue()

        def stdout_callback(line):
            if line.strip():
                output_queue.put_nowait(('message', line.strip()))

        def stderr_callback(line):
            if line.strip():
                if 'WARNING:' in line or 'DEPRECATION:' in line or 'A new release of pip is available' in line:
                    output_queue.put_nowait(('warning', line.strip()))
                else:
                    output_queue.put_nowait(('warning', line.strip()))
        run_task = asyncio.create_task(run_command_with_streaming(cmd, stdout_callback, stderr_callback, cwd, env))
        while not run_task.done() or not output_queue.empty():
            try:
                event_type, data = await asyncio.wait_for(output_queue.get(), timeout=0.1)
                yield {'event': event_type, 'data': data}
            except asyncio.TimeoutError:
                await asyncio.sleep(0.01)
        exit_code = await run_task
        if exit_code == 0:
            yield {'event': 'complete', 'data': 'Command completed successfully'}
        else:
            yield {'event': 'error', 'data': f'Command failed with exit code {exit_code}'}