tracer-cloud--opensre/surfaces/interactive_shell/ui/streaming/console.py

"""Rich console adapter for REPL dispatch streaming and cancellation."""

from __future__ import annotations

import sys
import threading
from typing import Any, Protocol

from rich.console import Console
from rich.file_proxy import FileProxy


class _PromptSpinner(Protocol):
    bytes_in: int
    streaming: bool

    def stop(self) -> None:
        raise NotImplementedError


class StreamingConsole(Console):
    """Console adapter for streaming progress + cancellation checks."""

    def __init__(
        self,
        spinner: _PromptSpinner,
        cancel_event: threading.Event,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self._spinner = spinner
        self._cancel_event = cancel_event

    def update_streaming_progress(self, bytes_received: int) -> None:
        self._spinner.bytes_in = bytes_received

    @property
    def cancel_requested(self) -> bool:
        return self._cancel_event.is_set()

    def print(self, *args: Any, **kwargs: Any) -> None:
        """Reset the TTY column before each print when not streaming."""
        if not self._spinner.streaming and not isinstance(sys.stdout, FileProxy):
            from surfaces.interactive_shell.ui.components.choice_menu import (
                ensure_tty_column_zero,
                prepare_repl_output_line,
            )
            from surfaces.interactive_shell.ui.components.rendering import (
                _repl_output_already_prepared,
                _repl_table_width,
            )

            if not args and not kwargs:
                ensure_tty_column_zero()
            elif not _repl_output_already_prepared():
                prepare_repl_output_line()
            if sys.stdout.isatty() and "width" not in kwargs:
                kwargs["width"] = _repl_table_width(self)
        super().print(*args, **kwargs)


__all__ = ["StreamingConsole"]