python-tty 0.1.0__py3-none-any.whl

python_tty/executor/executor.py

@@ -0,0 +1,335 @@
+import asyncio
+import inspect
+import time
+import uuid
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional
+
+from python_tty.config import ExecutorConfig
+from python_tty.ui.events import UIEvent, UIEventLevel
+from python_tty.ui.output import get_output_router
+from python_tty.executor.models import Invocation, RunState, RunStatus
+
+
+@dataclass
+class WorkItem:
+    invocation: Invocation
+    handler: Callable[[Invocation], object]
+
+
+class CommandExecutor:
+    def __init__(self, workers: int = 1, loop=None, config: ExecutorConfig = None):
+        if config is None:
+            config = ExecutorConfig(workers=workers)
+        self._config = config
+        self._worker_count = config.workers
+        self._loop = loop
+        self._queue = None
+        self._workers = []
+        self._locks: Dict[str, asyncio.Lock] = {}
+        self._runs: Dict[str, RunState] = {}
+        self._event_queues: Dict[str, asyncio.Queue] = {}
+        self._run_futures: Dict[str, asyncio.Future] = {}
+        self._retain_last_n = config.retain_last_n
+        self._ttl_seconds = config.ttl_seconds
+        self._pop_on_wait = config.pop_on_wait
+        self._output_router = get_output_router()
+
+    @property
+    def runs(self):
+        return self._runs
+
+    def start(self, loop=None):
+        if loop is not None:
+            self._loop = loop
+        if self._loop is None:
+            try:
+                self._loop = asyncio.get_running_loop()
+            except RuntimeError as exc:
+                raise RuntimeError("Executor start requires a running event loop") from exc
+        if self._queue is None:
+            self._queue = asyncio.Queue()
+        if self._workers:
+            return
+        for _ in range(self._worker_count):
+            self._workers.append(self._loop.create_task(self._worker_loop()))
+
+    def submit(self, invocation: Invocation, handler: Optional[Callable[[Invocation], object]] = None) -> str:
+        if invocation.run_id is None:
+            invocation.run_id = str(uuid.uuid4())
+        if handler is None:
+            handler = self._missing_handler
+        run_id = invocation.run_id
+        self._runs[run_id] = RunState(run_id=run_id)
+        if self._loop is None:
+            try:
+                self._loop = asyncio.get_running_loop()
+            except RuntimeError:
+                self._loop = None
+        if self._loop is None or not self._loop.is_running():
+            self._run_inline(invocation, handler)
+            return run_id
+        self.start()
+        self._run_futures[run_id] = self._loop.create_future()
+        self._queue.put_nowait(WorkItem(invocation=invocation, handler=handler))
+        return run_id
+
+    async def wait_result(self, run_id: str):
+        try:
+            future = self._run_futures.get(run_id)
+            if future is None:
+                run_state = self._runs.get(run_id)
+                if run_state is None:
+                    return None
+                if run_state.error is not None:
+                    raise run_state.error
+                return run_state.result
+            return await future
+        finally:
+            if self._pop_on_wait:
+                self.pop_run(run_id)
+
+    def wait_result_sync(self, run_id: str, timeout: Optional[float] = None):
+        try:
+            try:
+                running_loop = asyncio.get_running_loop()
+            except RuntimeError:
+                running_loop = None
+            if running_loop is not None and running_loop == self._loop:
+                raise RuntimeError("wait_result_sync cannot be called from the executor loop thread")
+            future = self._run_futures.get(run_id)
+            if future is not None and self._loop is not None and self._loop.is_running():
+                result_future = asyncio.run_coroutine_threadsafe(self.wait_result(run_id), self._loop)
+                return result_future.result(timeout)
+            run_state = self._runs.get(run_id)
+            if run_state is None:
+                return None
+            if run_state.status in (RunStatus.PENDING, RunStatus.RUNNING):
+                if self._loop is not None and self._loop.is_running():
+                    result_future = asyncio.run_coroutine_threadsafe(self.wait_result(run_id), self._loop)
+                    return result_future.result(timeout)
+                raise RuntimeError("Run is still pending but executor loop is not running")
+            if run_state.error is not None:
+                raise run_state.error
+            return run_state.result
+        finally:
+            if self._pop_on_wait:
+                self.pop_run(run_id)
+
+    def stream_events(self, run_id: str):
+        if self._loop is None or not self._loop.is_running():
+            raise RuntimeError("Event loop is not running")
+        return self._ensure_event_queue(run_id)
+
+    def publish_event(self, run_id: str, event):
+        if getattr(event, "run_id", None) is None:
+            event.run_id = run_id
+        if self._loop is not None and self._loop.is_running():
+            queue = self._ensure_event_queue(run_id)
+            self._queue_event(queue, event)
+        if self._output_router is not None:
+            self._output_router.emit(event)
+
+    async def shutdown(self, wait: bool = True):
+        workers = list(self._workers)
+        self._workers.clear()
+        for task in workers:
+            task.cancel()
+        if wait and workers:
+            await asyncio.gather(*workers, return_exceptions=True)
+
+    def shutdown_threadsafe(self, wait: bool = True, timeout: Optional[float] = None):
+        loop = self._loop
+        if loop is None or not loop.is_running():
+            for task in list(self._workers):
+                task.cancel()
+            self._workers.clear()
+            return None
+        future = asyncio.run_coroutine_threadsafe(self.shutdown(wait=wait), loop)
+        return future.result(timeout)
+
+    def submit_threadsafe(self, invocation: Invocation,
+                          handler: Optional[Callable[[Invocation], object]] = None) -> str:
+        if invocation.run_id is None:
+            invocation.run_id = str(uuid.uuid4())
+        if handler is None:
+            handler = self._missing_handler
+        run_id = invocation.run_id
+        self._runs[run_id] = RunState(run_id=run_id)
+        if self._loop is None or not self._loop.is_running():
+            return self.submit(invocation, handler=handler)
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+        if running_loop == self._loop:
+            return self.submit(invocation, handler=handler)
+
+        async def _enqueue():
+            self.start()
+            self._run_futures[run_id] = self._loop.create_future()
+            self._queue.put_nowait(WorkItem(invocation=invocation, handler=handler))
+
+        asyncio.run_coroutine_threadsafe(_enqueue(), self._loop).result()
+        return run_id
+
+    async def _worker_loop(self):
+        while True:
+            work_item = await self._queue.get()
+            run_state = self._runs.get(work_item.invocation.run_id)
+            lock = self._locks.setdefault(work_item.invocation.lock_key, asyncio.Lock())
+            async with lock:
+                await self._execute_work_item(work_item, run_state)
+            self._queue.task_done()
+
+    async def _execute_work_item(self, work_item: WorkItem, run_state: Optional[RunState]):
+        if run_state is None:
+            return
+        run_state.status = RunStatus.RUNNING
+        run_state.started_at = time.time()
+        self.publish_event(run_state.run_id, self._build_run_event("start", UIEventLevel.INFO))
+        try:
+            result = work_item.handler(work_item.invocation)
+            if inspect.isawaitable(result):
+                result = await result
+            run_state.result = result
+            run_state.status = RunStatus.SUCCEEDED
+            self.publish_event(run_state.run_id, self._build_run_event("success", UIEventLevel.SUCCESS))
+            self._resolve_future(run_state, result=result)
+        except Exception as exc:
+            run_state.error = exc
+            run_state.status = RunStatus.FAILED
+            self.publish_event(
+                run_state.run_id,
+                self._build_run_event("failure", UIEventLevel.ERROR, payload={"error": str(exc)}),
+            )
+            self._resolve_future(run_state, error=exc)
+        finally:
+            run_state.finished_at = time.time()
+            self._cleanup_runs()
+
+    def _resolve_future(self, run_state: RunState, result=None, error: Optional[BaseException] = None):
+        future = self._run_futures.get(run_state.run_id)
+        if future is None or future.done():
+            return
+        if error is not None:
+            future.set_exception(error)
+        else:
+            future.set_result(result)
+
+    def _run_inline(self, invocation: Invocation, handler):
+        run_state = self._runs.get(invocation.run_id)
+        if run_state is None:
+            return
+        run_state.status = RunStatus.RUNNING
+        run_state.started_at = time.time()
+        self.publish_event(run_state.run_id, self._build_run_event("start", UIEventLevel.INFO))
+        try:
+            result = handler(invocation)
+            if inspect.isawaitable(result):
+                result = self._run_awaitable_inline(result)
+            run_state.result = result
+            run_state.status = RunStatus.SUCCEEDED
+            self.publish_event(run_state.run_id, self._build_run_event("success", UIEventLevel.SUCCESS))
+        except Exception as exc:
+            run_state.error = exc
+            run_state.status = RunStatus.FAILED
+            self.publish_event(
+                run_state.run_id,
+                self._build_run_event("failure", UIEventLevel.ERROR, payload={"error": str(exc)}),
+            )
+        finally:
+            run_state.finished_at = time.time()
+            self._cleanup_runs()
+
+    @staticmethod
+    def _build_run_event(event_type: str, level: UIEventLevel, payload=None):
+        return UIEvent(msg=event_type, level=level, event_type=event_type, payload=payload)
+
+    @staticmethod
+    def _missing_handler(invocation: Invocation):
+        raise RuntimeError("No handler provided for invocation execution")
+
+    def _ensure_event_queue(self, run_id: str):
+        if self._loop is None or not self._loop.is_running():
+            return self._event_queues.setdefault(run_id, asyncio.Queue())
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+        if running_loop == self._loop:
+            return self._event_queues.setdefault(run_id, asyncio.Queue())
+        future = asyncio.run_coroutine_threadsafe(self._create_event_queue(run_id), self._loop)
+        return future.result()
+
+    async def _create_event_queue(self, run_id: str):
+        return self._event_queues.setdefault(run_id, asyncio.Queue())
+
+    def _queue_event(self, queue: asyncio.Queue, event):
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+        if running_loop == self._loop:
+            queue.put_nowait(event)
+        else:
+            self._loop.call_soon_threadsafe(queue.put_nowait, event)
+
+    def _run_awaitable_inline(self, awaitable):
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+        if asyncio.isfuture(awaitable):
+            raise RuntimeError("Inline awaitable must be a coroutine, not a Future/Task")
+        if running_loop is None:
+            return asyncio.run(self._awaitable_to_coroutine(awaitable))
+        if self._loop is not None and self._loop.is_running():
+            future = asyncio.run_coroutine_threadsafe(
+                self._awaitable_to_coroutine(awaitable),
+                self._loop,
+            )
+            return future.result()
+        raise RuntimeError("Cannot run awaitable inline while an event loop is running")
+
+    @staticmethod
+    def _awaitable_to_coroutine(awaitable):
+        if asyncio.iscoroutine(awaitable):
+            return awaitable
+
+        async def _await_obj():
+            return await awaitable
+
+        return _await_obj()
+
+    def pop_run(self, run_id: str):
+        run_state = self._runs.pop(run_id, None)
+        future = self._run_futures.pop(run_id, None)
+        if future is not None and not future.done():
+            future.cancel()
+        self._event_queues.pop(run_id, None)
+        return run_state
+
+    def _cleanup_runs(self):
+        if self._retain_last_n is None and self._ttl_seconds is None:
+            return
+        now = time.time()
+        completed = []
+        for run_id, run_state in self._runs.items():
+            if run_state.status in (RunStatus.PENDING, RunStatus.RUNNING):
+                continue
+            completed.append((run_state.finished_at, run_id))
+        remove_ids = set()
+        if self._ttl_seconds is not None:
+            for finished_at, run_id in completed:
+                if finished_at is None:
+                    continue
+                if now - finished_at >= self._ttl_seconds:
+                    remove_ids.add(run_id)
+        if self._retain_last_n is not None and self._retain_last_n >= 0:
+            completed.sort(reverse=True)
+            for _, run_id in completed[self._retain_last_n:]:
+                remove_ids.add(run_id)
+        for run_id in remove_ids:
+            self.pop_run(run_id)
+

python_tty/executor/models.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+
+class RunStatus(Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    TIMEOUT = "timeout"
+
+
+@dataclass
+class Invocation:
+    run_id: Optional[str] = None
+    source: str = "tty"
+    principal: Optional[str] = None
+    console_id: Optional[str] = None
+    command_id: Optional[str] = None
+    command_name: Optional[str] = None
+    argv: List[str] = field(default_factory=list)
+    kwargs: Dict[str, Any] = field(default_factory=dict)
+    lock_key: str = "global"
+    timeout_ms: Optional[int] = None
+    audit_policy: Optional[str] = None
+    raw_cmd: Optional[str] = None
+
+
+@dataclass
+class RunState:
+    run_id: str
+    status: RunStatus = RunStatus.PENDING
+    result: Any = None
+    error: Optional[BaseException] = None
+    started_at: Optional[float] = None
+    finished_at: Optional[float] = None

python_tty/ui/__init__.py

@@ -0,0 +1,13 @@
+from python_tty.ui.events import UIEvent, UIEventLevel, UIEventListener, UIEventSpeaker
+from python_tty.ui.output import OutputRouter, get_output_router, proxy_print
+
+__all__ = [
+    "UIEvent",
+    "UIEventLevel",
+    "UIEventListener",
+    "UIEventSpeaker",
+    "OutputRouter",
+    "get_output_router",
+    "proxy_print",
+]
+

python_tty/ui/events.py

@@ -0,0 +1,55 @@
+import enum
+
+
+class UIEventLevel(enum.Enum):
+    TEXT = -1
+    INFO = 0
+    WARNING = 1
+    ERROR = 2
+    SUCCESS = 3
+    FAILURE = 4
+    DEBUG = 5
+
+    @staticmethod
+    def map_level(code):
+        if code == 0:
+            return UIEventLevel.INFO
+        elif code == 1:
+            return UIEventLevel.WARNING
+        elif code == 2:
+            return UIEventLevel.ERROR
+        elif code == 3:
+            return UIEventLevel.SUCCESS
+        elif code == 4:
+            return UIEventLevel.FAILURE
+        elif code == 5:
+            return UIEventLevel.DEBUG
+
+
+class UIEvent:
+    def __init__(self, msg, level=UIEventLevel.TEXT, run_id=None, event_type=None, payload=None):
+        self.msg = msg
+        self.level = level
+        self.run_id = run_id
+        self.event_type = event_type
+        self.payload = payload
+
+
+class UIEventListener:
+    def handler_event(self, event: UIEvent):
+        pass
+
+
+class UIEventSpeaker:
+    def __init__(self):
+        self._event_listener = []
+
+    def add_event_listener(self, listener: UIEventListener):
+        self._event_listener.append(listener)
+
+    def remove_event_listener(self, listener: UIEventListener):
+        self._event_listener.remove(listener)
+
+    def notify_event_listeners(self, event: UIEvent):
+        for listener in self._event_listener:
+            listener.handler_event(event)

python_tty/ui/output.py

@@ -0,0 +1,102 @@
+import threading
+
+from prompt_toolkit import print_formatted_text
+from prompt_toolkit.formatted_text import FormattedText
+from prompt_toolkit.styles import Style
+
+from python_tty.ui.events import UIEvent, UIEventLevel
+
+
+MSG_LEVEL_SYMBOL = {
+    0: "[*] ",
+    1: "[!] ",
+    2: "[x] ",
+    3: "[+] ",
+    4: "[-] ",
+    5: "[@] "
+}
+
+MSG_LEVEL_SYMBOL_STYLE = {
+    0: "fg:green",
+    1: "fg:yellow",
+    2: "fg:red",
+    3: "fg:blue",
+    4: "fg:white",
+    5: "fg:pink"
+}
+
+
+class OutputRouter:
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._app = None
+        self._output = None
+
+    def bind_session(self, session):
+        if session is None:
+            return
+        with self._lock:
+            self._app = getattr(session, "app", None)
+            self._output = getattr(session, "output", None)
+
+    def clear_session(self, session=None):
+        with self._lock:
+            if session is None or getattr(session, "app", None) == self._app:
+                self._app = None
+                self._output = None
+
+    def emit(self, event: UIEvent):
+        with self._lock:
+            app = self._app
+            output = self._output
+
+        def _render():
+            text, style = _format_event(event)
+            if output is not None:
+                print_formatted_text(text, style=style, output=output)
+            else:
+                print_formatted_text(text, style=style)
+
+        if app is not None and getattr(app, "is_running", False):
+            if hasattr(app, "call_from_executor") and hasattr(app, "run_in_terminal"):
+                app.call_from_executor(lambda: app.run_in_terminal(_render))
+                return
+        _render()
+
+
+def _normalize_level(level):
+    if isinstance(level, UIEventLevel):
+        return level
+    if level is None:
+        return UIEventLevel.TEXT
+    if level == UIEventLevel.TEXT.value:
+        return UIEventLevel.TEXT
+    mapped = UIEventLevel.map_level(level)
+    return UIEventLevel.TEXT if mapped is None else mapped
+
+
+def _format_event(event: UIEvent):
+    level = _normalize_level(event.level)
+    if level == UIEventLevel.TEXT:
+        return event.msg, None
+    formatted_text = FormattedText([
+        ("class:level", MSG_LEVEL_SYMBOL[level.value]),
+        ("class:text", str(event.msg)),
+    ])
+    style = Style.from_dict({
+        "level": MSG_LEVEL_SYMBOL_STYLE[level.value]
+    })
+    return formatted_text, style
+
+
+_OUTPUT_ROUTER = OutputRouter()
+
+
+def get_output_router() -> OutputRouter:
+    return _OUTPUT_ROUTER
+
+
+def proxy_print(text="", text_type=UIEventLevel.TEXT):
+    event = UIEvent(msg=text, level=_normalize_level(text_type))
+    get_output_router().emit(event)
+

python_tty/utils/__init__.py

@@ -0,0 +1,13 @@
+from python_tty.utils.table import Table
+from python_tty.utils.tokenize import get_command_token, get_func_param_strs, split_cmd, tokenize_cmd
+from python_tty.utils.ui_logger import ConsoleHandler
+
+__all__ = [
+    "ConsoleHandler",
+    "Table",
+    "get_command_token",
+    "get_func_param_strs",
+    "split_cmd",
+    "tokenize_cmd",
+]
+

python_tty/utils/table.py

@@ -0,0 +1,126 @@
+import copy
+
+
+class Cell:
+    def __init__(self, data):
+        self.data = data
+        self.data_str = str(self.data)
+        self.data_width = len(self.data_str)
+        self.padding = ""
+
+    def update_max_width(self, padding_len: int):
+        if padding_len > 0:
+            self.padding = " " * padding_len
+
+    def __str__(self):
+        return "".join([self.data_str, self.padding])
+
+
+class HeaderCell(Cell):
+    def __init__(self, data, seq="-"):
+        super().__init__(data)
+        self.seq_str = self.data_width * seq
+        self.data_str = str(self.data)
+
+    def get_seq_str(self):
+        return "".join([self.seq_str, self.padding])
+
+
+class Table:
+    def __init__(self, header: [], data: [[]], title="",
+                 title_indent=0, data_indent=4, data_seq_len=4,
+                 title_seq="=", header_seq="-", header_footer=True):
+        self.title = title
+        self.title_indent = title_indent
+        self.data_indent = data_indent
+        self.data_seq_len = data_seq_len
+        self.title_seq = title_seq
+        self.header_seq = header_seq
+        self.header_footer = header_footer
+        self.header = copy.deepcopy(header)
+        self.data = copy.deepcopy(data)
+        self._padding_data()
+        self._format_header()
+        self._merge_data()
+        self._padding_max_width()
+
+    def _padding_data(self):
+        table_header_item_num = len(self.header)
+        for i in range(len(self.data)):
+            row = self.data[i]
+            if len(row) < table_header_item_num:
+                self.data[i].append("")
+            elif len(row) > table_header_item_num:
+                self.data[i] = row[:table_header_item_num]
+
+    def _format_header(self):
+        for i in range(len(self.header)):
+            cell = self.header[i]
+            self.header[i] = str(cell)[0:1].upper() + str(cell)[1:]
+
+    def _merge_data(self):
+        data = []
+        header = []
+        for cell in self.header:
+            header.append(HeaderCell(cell, self.header_seq))
+        data.append(header)
+        for row in self.data:
+            line = []
+            for cell in row:
+                line.append(Cell(cell))
+            data.append(line)
+        self.data = data
+
+    def _padding_max_width(self):
+        max_widths = [len(cell) for cell in self.header]
+        for i in range(len(self.data)):
+            for j in range(len(self.data[i])):
+                max_width = max_widths[j]
+                cell = self.data[i][j]
+                if cell.data_width > max_width:
+                    max_widths[j] = cell.data_width
+        for i in range(len(self.data)):
+            for j in range(len(self.data[i])):
+                max_width = max_widths[j]
+                cell = self.data[i][j]
+                if cell.data_width < max_width:
+                    cell.update_max_width(max_width - cell.data_width)
+
+    def print_row(self, row: []):
+        cells = []
+        seqs = []
+        for cell in row:
+            if isinstance(cell, HeaderCell):
+                seqs.append(cell.get_seq_str())
+            cells.append(str(cell))
+        if len(seqs) > 0:
+            return " "*self.data_indent + (" " * self.data_seq_len).join(cells),\
+                " "*self.data_indent + (" " * self.data_seq_len).join(seqs)
+        else:
+            return " "*self.data_indent + (" " * self.data_seq_len).join(cells), None
+
+    def print_data(self):
+        lines = []
+        for row in self.data:
+            line, seq = self.print_row(row)
+            lines.append(line)
+            if seq is not None:
+                lines.append(seq)
+        return "\n".join(lines)
+
+    def print_title(self):
+        title_str = str(self.title)
+        if title_str != "":
+            if not title_str[0:1].isupper():
+                title_str = title_str[0:1].upper() + title_str[1:].lower()
+            title_line = " "*self.title_indent + title_str
+            seq_line = " "*self.title_indent + self.title_seq*len(self.title)
+            return "\n".join([title_line, seq_line])
+
+    def __str__(self):
+        if str(self.title) != "":
+            title = self.print_title() + "\n\n"
+            table_str = title + self.print_data()
+        else:
+            table_str = self.print_data()
+        return "\n" + table_str + "\n" if self.header_footer else table_str