stata-code 0.3.0__py3-none-any.whl

stata_code/core/_refs.py

@@ -0,0 +1,97 @@
+"""Process-wide ref store for log/graph/matrix payloads.
+
+`stata_code` returns large payloads (full logs, graph bytes, matrix values)
+by reference rather than inline, to keep agent token economy in check. This
+module owns the underlying dict; auxiliary tools (`get_log`, `get_graph`,
+`get_matrix`) read from it.
+
+Refs are valid only within the lifetime of the producing process. Per
+SCHEMA.md §3.3: consumers MUST NOT persist refs across sessions.
+
+Eviction: the store is bounded by `MAX_ENTRIES` (default 256) on an LRU
+basis. When inserting beyond the cap, the least-recently-used entry is
+dropped. This guards long-running MCP server processes from unbounded
+growth. Adjust via `set_capacity(n)` if you need more slack.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections import OrderedDict
+from typing import Any
+
+DEFAULT_MAX_ENTRIES = 256
+
+_lock = threading.Lock()
+_store: OrderedDict[str, Any] = OrderedDict()
+_max_entries: int = DEFAULT_MAX_ENTRIES
+
+
+def set_capacity(n: int) -> None:
+    """Set the maximum number of refs to retain. Must be ≥ 1."""
+    global _max_entries
+    if n < 1:
+        raise ValueError("capacity must be ≥ 1")
+    with _lock:
+        _max_entries = n
+        _evict_to_capacity_locked()
+
+
+def get_capacity() -> int:
+    with _lock:
+        return _max_entries
+
+
+def put(ref: str, payload: Any) -> None:
+    with _lock:
+        # If already present, move-to-end keeps the entry "fresh".
+        if ref in _store:
+            _store.move_to_end(ref)
+        _store[ref] = payload
+        _evict_to_capacity_locked()
+
+
+def get(ref: str) -> Any | None:
+    """Fetch a payload. Touches LRU order so frequently-used refs survive."""
+    with _lock:
+        if ref not in _store:
+            return None
+        _store.move_to_end(ref)
+        return _store[ref]
+
+
+def has(ref: str) -> bool:
+    with _lock:
+        return ref in _store
+
+
+def discard(ref: str) -> None:
+    with _lock:
+        _store.pop(ref, None)
+
+
+def clear_prefix(prefix: str) -> int:
+    """Drop all refs whose key starts with `prefix`. Returns count dropped."""
+    with _lock:
+        keys = [k for k in _store if k.startswith(prefix)]
+        for k in keys:
+            del _store[k]
+        return len(keys)
+
+
+def clear_all() -> None:
+    """Drop every ref. Mainly for tests."""
+    with _lock:
+        _store.clear()
+
+
+def size() -> int:
+    with _lock:
+        return len(_store)
+
+
+def _evict_to_capacity_locked() -> None:
+    """Drop oldest entries until len(_store) <= _max_entries. Caller holds _lock."""
+    while len(_store) > _max_entries:
+        # popitem(last=False) drops the OLDEST (FIFO end of the OrderedDict)
+        _store.popitem(last=False)

stata_code/core/_runtime.py

@@ -0,0 +1,179 @@
+"""Private pystata runtime wrapper — process-wide singleton.
+
+pystata initializes one Stata kernel per Python process. This module manages
+that lifecycle and exposes a small surface (`run_capture`, sfi accessors) for
+the public runner. Not part of the API; subject to change without notice.
+"""
+
+from __future__ import annotations
+
+import io
+import re
+import sys
+import threading
+from contextlib import redirect_stdout
+from pathlib import Path
+from typing import Any
+
+# Common pystata install locations across platforms. Order is best-guess for
+# popularity; the first match wins. Users can also pre-install pystata into
+# their Python environment, in which case we skip the path search entirely.
+_PYSTATA_SEARCH_PATHS: tuple[str, ...] = (
+    "/Applications/Stata/utilities",
+    "/Applications/Stata18/utilities",
+    "/Applications/Stata17/utilities",
+    "/usr/local/stata18/utilities",
+    "/usr/local/stata17/utilities",
+    r"C:\Program Files\Stata18\utilities",
+    r"C:\Program Files\Stata17\utilities",
+)
+
+_RC_RE = re.compile(r"r\((-?\d+)\);")
+
+_lock = threading.Lock()
+_runtime: PystataRuntime | None = None
+_init_attempted = False
+_init_error: Exception | None = None
+
+
+class PystataNotAvailable(RuntimeError):
+    """Raised when pystata cannot be imported or Stata cannot be initialized."""
+
+
+class PystataRuntime:
+    """Wraps a single pystata-initialized Stata kernel.
+
+    Construct via `get_runtime()`; do not instantiate directly. The kernel
+    persists for the lifetime of the Python process.
+    """
+
+    def __init__(self) -> None:
+        self._initialized = False
+        self._stata: Any = None
+        self._config: Any = None
+        self._sfi: Any = None
+        self._edition: str = ""
+
+    @staticmethod
+    def _find_pystata_path() -> str | None:
+        try:
+            import pystata  # noqa: F401
+
+            return None  # already importable; nothing to add to sys.path
+        except ImportError:
+            pass
+        for p in _PYSTATA_SEARCH_PATHS:
+            if Path(p).joinpath("pystata").is_dir():
+                return p
+        return None
+
+    def _initialize(self) -> None:
+        if self._initialized:
+            return
+        path = self._find_pystata_path()
+        if path is not None and path not in sys.path:
+            sys.path.insert(0, path)
+        try:
+            from pystata import config as cfg
+        except ImportError as exc:
+            raise PystataNotAvailable(
+                "pystata is not importable; install Stata 17+ or set PYTHONPATH"
+            ) from exc
+
+        last_err: Exception | None = None
+        for ed in ("mp", "se", "be"):
+            try:
+                cfg.init(ed, splash=False)
+                self._edition = ed
+                break
+            except Exception as exc:  # noqa: BLE001
+                last_err = exc
+        else:
+            raise PystataNotAvailable(
+                f"pystata.config.init failed for all editions: {last_err}"
+            )
+
+        cfg.set_graph_show(False)
+        cfg.set_graph_format("png")
+
+        import sfi  # type: ignore[import-not-found]
+        from pystata import stata as st  # must be imported after config.init
+
+        self._config = cfg
+        self._stata = st
+        self._sfi = sfi
+        self._initialized = True
+
+    @property
+    def is_initialized(self) -> bool:
+        return self._initialized
+
+    @property
+    def edition(self) -> str:
+        return self._edition
+
+    @property
+    def stata(self) -> Any:
+        return self._stata
+
+    @property
+    def sfi(self) -> Any:
+        return self._sfi
+
+    def run_capture(self, code: str) -> tuple[str, int, str | None]:
+        """Run Stata code and return (captured_stdout, rc, error_message).
+
+        - rc == 0 on success.
+        - rc > 0 is Stata's `_rc` (parsed from the SystemError message).
+        - rc == -1 is an adapter-level crash (non-Stata exception).
+        - error_message is None on success, non-None on failure.
+        """
+        buf = io.StringIO()
+        try:
+            with redirect_stdout(buf):
+                self._stata.run(code, quietly=False, echo=False)
+            return buf.getvalue(), 0, None
+        except SystemError as exc:
+            err_text = str(exc.args[0]) if exc.args else str(exc)
+            m = _RC_RE.search(err_text)
+            rc = int(m.group(1)) if m else -1
+            return buf.getvalue(), rc, err_text.strip()
+        except Exception as exc:  # noqa: BLE001
+            return buf.getvalue(), -1, f"{type(exc).__name__}: {exc}"
+
+
+def get_runtime() -> PystataRuntime:
+    """Return the process-wide runtime, initializing on first call.
+
+    Raises `PystataNotAvailable` if pystata or Stata cannot be initialized.
+    The failure is cached: subsequent calls re-raise the same error without
+    retrying init (which would print messy errors repeatedly).
+    """
+    global _runtime, _init_attempted, _init_error
+    if _runtime is not None:
+        return _runtime
+    if _init_attempted and _init_error is not None:
+        raise _init_error  # cached failure
+    with _lock:
+        if _runtime is not None:
+            return _runtime
+        _init_attempted = True
+        try:
+            rt = PystataRuntime()
+            rt._initialize()
+            _runtime = rt
+            return _runtime
+        except Exception as exc:  # noqa: BLE001
+            _init_error = exc if isinstance(exc, PystataNotAvailable) else (
+                PystataNotAvailable(f"runtime init failed: {exc}")
+            )
+            raise _init_error from exc
+
+
+def is_available() -> bool:
+    """Cheap, non-raising check. Initializes on first call."""
+    try:
+        get_runtime()
+        return True
+    except PystataNotAvailable:
+        return False