mindroot.lib.streamcmd - MindRoot 9.10.0 documentation

Hide navigation sidebar
Hide table of contents sidebar
Skip to content
Back to top
Toggle table of contents sidebar

# Source code for mindroot.lib.streamcmd

    #!/usr/bin/env python3

    import asyncio
    import sys
    from typing import Optional, Callable, Any, List, Dict, AsyncGenerator

    # ANSI color codes for terminal output
    COLORS = {
        'reset': '\033[0m',
        'green': '\033[32m',  # stdout
        'yellow': '\033[33m',  # stderr/warning
        'blue': '\033[34m',    # info
        'red': '\033[31m'      # error
    }

    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
            )

            # Create tasks to read from stdout and stderr
            stdout_task = asyncio.create_task(
                read_stream(process.stdout, stdout_callback)
            )
            stderr_task = asyncio.create_task(
                read_stream(process.stderr, stderr_callback)
            )

            # Wait for the process to complete and streams to be fully read
            await asyncio.gather(stdout_task, stderr_task)
            exit_code = await process.wait()

            return exit_code

        except Exception as e:
            print(f"Failed to run command: {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
        """
        # Debug output
        print(f"{COLORS['blue']}[DEBUG] Running command: {' '.join(cmd)}{COLORS['reset']}")

        # Send initial event
        yield {"event": "message", "data": f"Running command: {' '.join(cmd)}"}

        # Create queues for stdout and stderr
        output_queue = asyncio.Queue()

        # Define callbacks for stdout and stderr
        def stdout_callback(line):
            if line.strip():
                print(f"{COLORS['green']}[STDOUT] {line.strip()}{COLORS['reset']}")
                output_queue.put_nowait(("message", line.strip()))

        def stderr_callback(line):
            if line.strip():
                # Determine if this is a warning or an error
                if ("WARNING:" in line or
                    "DEPRECATION:" in line or
                    "A new release of pip is available" in line):
                    print(f"{COLORS['yellow']}[WARNING] {line.strip()}{COLORS['reset']}")
                    output_queue.put_nowait(("warning", line.strip()))
                else:
                    print(f"{COLORS['red']}[ERROR] {line.strip()}{COLORS['reset']}")
                    output_queue.put_nowait(("warning", line.strip()))

        # Run the command in a separate task
        run_task = asyncio.create_task(
            run_command_with_streaming(cmd, stdout_callback, stderr_callback, cwd, env)
        )

        # Stream events from the queue while the command is running
        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:
                # No output available, just continue
                await asyncio.sleep(0.01)

        # Get the exit code
        exit_code = await run_task

        # Send completion event
        print(f"{COLORS['blue']}[INFO] Command completed with exit code {exit_code}{COLORS['reset']}")
        if exit_code == 0:
            yield {"event": "complete", "data": "Command completed successfully"}
        else:
            print(f"{COLORS['red']}[ERROR] Command failed with exit code {exit_code}{COLORS['reset']}")
            yield {"event": "error", "data": f"Command failed with exit code {exit_code}"}