mcp-stata 1.0.1__py3-none-any.whl

mcp_stata/stata_client.py

@@ -0,0 +1,524 @@
+import sys
+import os
+import json
+import re
+import base64
+import logging
+import threading
+import time
+from io import StringIO
+from contextlib import contextmanager
+from typing import Any, List, Optional, Dict
+import pandas as pd
+from .discovery import find_stata_path
+from .smcl.smcl2html import smcl_to_markdown
+from .models import (
+    CommandResponse,
+    ErrorEnvelope,
+    GraphExport,
+    GraphExportResponse,
+    GraphInfo,
+    GraphListResponse,
+    VariableInfo,
+    VariablesResponse,
+)
+
+logger = logging.getLogger("mcp_stata")
+
+class StataClient:
+    _instance = None
+    _initialized = False
+    _exec_lock: threading.Lock
+    MAX_DATA_ROWS = 500
+    MAX_GRAPH_BYTES = 50 * 1024 * 1024  # Allow large graph exports (~50MB)
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(StataClient, cls).__new__(cls)
+            cls._instance._exec_lock = threading.Lock()
+        return cls._instance
+
+    @contextmanager
+    def _redirect_io(self):
+        """Safely redirect stdout/stderr for the duration of a Stata call."""
+        out_buf, err_buf = StringIO(), StringIO()
+        backup_stdout, backup_stderr = sys.stdout, sys.stderr
+        sys.stdout, sys.stderr = out_buf, err_buf
+        try:
+            yield out_buf, err_buf
+        finally:
+            sys.stdout, sys.stderr = backup_stdout, backup_stderr
+
+    def init(self):
+        """Initializes usage of pystata."""
+        if self._initialized:
+            return
+
+        try:
+            # 1. Setup config
+            # 1. Setup config
+            import stata_setup
+            try:
+                stata_exec_path, edition = find_stata_path()
+            except FileNotFoundError as e:
+                raise RuntimeError(f"Stata binary not found: {e}") from e
+            except PermissionError as e:
+                raise RuntimeError(
+                    f"Stata binary is not executable: {e}. "
+                    "Point STATA_PATH directly to the Stata binary (e.g., .../Contents/MacOS/stata-mp)."
+                ) from e
+            logger.info(f"Discovery found Stata at: {stata_exec_path} ({edition})")
+            
+            # Helper to try init
+            def tries_init(path_to_try):
+                try:
+                    logger.info(f"Attempting stata_setup.config with: {path_to_try}")
+                    stata_setup.config(path_to_try, edition)
+                    return True
+                except Exception as e:
+                    logger.warning(f"Init failed with {path_to_try}: {e}")
+                    return False
+
+            success = False
+            candidates = []
+            
+            # 1. Binary Dir: .../Contents/MacOS
+            bin_dir = os.path.dirname(stata_exec_path)
+            
+            # 2. App Bundle: .../StataMP.app
+            # Walk up to find .app
+            curr = bin_dir
+            app_bundle = None
+            while len(curr) > 1:
+                if curr.endswith(".app"):
+                    app_bundle = curr
+                    break
+                curr = os.path.dirname(curr)
+                
+            if app_bundle:
+                # Priority 1: The installation root (parent of .app)
+                candidates.append(os.path.dirname(app_bundle))
+                
+                # Priority 2: The .app bundle itself
+                candidates.append(app_bundle)
+            
+            # Priority 3: The binary directory
+            candidates.append(bin_dir)
+            
+            for path in candidates:
+                if tries_init(path):
+                    success = True
+                    break
+            
+            if not success:
+                raise RuntimeError(
+                    f"stata_setup.config failed. Tried: {candidates}. "
+                    f"Derived from binary: {stata_exec_path}"
+                )
+            
+            # 2. Import pystata
+            from pystata import stata
+            self.stata = stata
+            self._initialized = True
+            
+        except ImportError:
+            # Fallback for when stata_setup isn't in PYTHONPATH yet?
+            # Usually users must have it installed. We rely on discovery logic.
+            raise RuntimeError("Could not import `stata_setup`. Ensure pystata is installed.")
+        except Exception as e:
+            raise RuntimeError(f"Failed to initialize Stata: {e}")
+
+    def _read_return_code(self) -> int:
+        """Read the last Stata return code without mutating rc."""
+        try:
+            from sfi import Macro
+            rc_val = Macro.getCValue("rc")  # type: ignore[attr-defined]
+            return int(float(rc_val))
+        except Exception:
+            try:
+                self.stata.run("global MCP_RC = c(rc)")
+                from sfi import Macro as Macro2
+                rc_val = Macro2.getGlobal("MCP_RC")
+                return int(float(rc_val))
+            except Exception:
+                return -1
+
+    def _parse_rc_from_text(self, text: str) -> Optional[int]:
+        match = re.search(r"r\((\d+)\)", text)
+        if match:
+            try:
+                return int(match.group(1))
+            except Exception:
+                return None
+        return None
+
+    def _parse_line_from_text(self, text: str) -> Optional[int]:
+        match = re.search(r"line\s+(\d+)", text, re.IGNORECASE)
+        if match:
+            try:
+                return int(match.group(1))
+            except Exception:
+                return None
+        return None
+
+    def _smcl_to_text(self, smcl: str) -> str:
+        """Convert simple SMCL markup into plain text for LLM-friendly help."""
+        # First, keep inline directive content if present (e.g., {bf:word} -> word)
+        cleaned = re.sub(r"\{[^}:]+:([^}]*)\}", r"\1", smcl)
+        # Remove remaining SMCL brace commands like {smcl}, {vieweralsosee ...}, {txt}, {p}
+        cleaned = re.sub(r"\{[^}]*\}", "", cleaned)
+        # Normalize whitespace
+        cleaned = cleaned.replace("\r", "")
+        lines = [line.rstrip() for line in cleaned.splitlines()]
+        return "\n".join(lines).strip()
+
+    def _build_error_envelope(
+        self,
+        command: str,
+        rc: int,
+        stdout: str,
+        stderr: str,
+        exc: Optional[Exception],
+        trace: bool,
+    ) -> ErrorEnvelope:
+        combined = "\n".join(filter(None, [stdout, stderr, str(exc) if exc else ""])).strip()
+        rc_hint = self._parse_rc_from_text(combined) if combined else None
+        rc_final = rc if rc not in (-1, None) else rc_hint
+        line_no = self._parse_line_from_text(combined) if combined else None
+        snippet = combined[-800:] if combined else None
+        message = (stderr or (str(exc) if exc else "") or stdout or "Stata error").strip()
+        return ErrorEnvelope(
+            message=message,
+            rc=rc_final,
+            line=line_no,
+            command=command,
+            stdout=stdout or None,
+            stderr=stderr or None,
+            snippet=snippet,
+            trace=trace or None,
+        )
+
+    def _exec_with_capture(self, code: str, echo: bool = True, trace: bool = False) -> CommandResponse:
+        """Execute Stata code with stdout/stderr capture and rc detection."""
+        if not self._initialized:
+            self.init()
+
+        start_time = time.time()
+        exc: Optional[Exception] = None
+        with self._exec_lock:
+            with self._redirect_io() as (out_buf, err_buf):
+                try:
+                    if trace:
+                        self.stata.run("set trace on")
+                    self.stata.run(code, echo=echo)
+                except Exception as e:
+                    exc = e
+                finally:
+                    rc = self._read_return_code()
+                    if trace:
+                        try:
+                            self.stata.run("set trace off")
+                        except Exception:
+                            pass
+
+        stdout = out_buf.getvalue()
+        stderr = err_buf.getvalue()
+        # If no exception and stderr is empty, treat rc anomalies as success (e.g., spurious rc reads)
+        if exc is None and (not stderr or not stderr.strip()):
+            rc = 0 if rc is None or rc != 0 else rc
+        success = rc == 0 and exc is None
+        error = None
+        if not success:
+            error = self._build_error_envelope(code, rc, stdout, stderr, exc, trace)
+        duration = time.time() - start_time
+        code_preview = code.replace("\n", "\\n")
+        logger.info(
+            "stata.run rc=%s success=%s trace=%s duration_ms=%.2f code_preview=%s",
+            rc,
+            success,
+            trace,
+            duration * 1000,
+            code_preview[:120],
+        )
+        return CommandResponse(
+            command=code,
+            rc=rc,
+            stdout=stdout,
+            stderr=stderr or None,
+            success=success,
+            error=error,
+        )
+
+    def run_command(self, code: str, echo: bool = True) -> str:
+        """Runs a Stata command and returns raw output (legacy)."""
+        result = self._exec_with_capture(code, echo=echo)
+        if result.success:
+            return result.stdout
+        if result.error:
+            return f"Error executing Stata code (r({result.error.rc})):\n{result.error.message}"
+        return result.stdout or "Unknown Stata error"
+
+    def run_command_structured(self, code: str, echo: bool = True, trace: bool = False) -> CommandResponse:
+        """Runs a Stata command and returns a structured envelope."""
+        return self._exec_with_capture(code, echo=echo, trace=trace)
+
+    def get_data(self, start: int = 0, count: int = 50) -> List[Dict[str, Any]]:
+        """Returns valid JSON-serializable data."""
+        if not self._initialized:
+            self.init()
+
+        if count > self.MAX_DATA_ROWS:
+            count = self.MAX_DATA_ROWS
+
+        try:
+            # Use pystata integration to retrieve data
+            df = self.stata.pdataframe_from_data()
+
+            # Slice
+            sliced = df.iloc[start : start + count]
+
+            # Convert to dict
+            return sliced.to_dict(orient="records")
+        except Exception as e:
+            return [{"error": f"Failed to retrieve data: {e}"}]
+
+    def list_variables(self) -> List[Dict[str, str]]:
+        """Returns list of variables with labels."""
+        if not self._initialized:
+            self.init()
+            
+        # We can use sfi to be efficient
+        from sfi import Data
+        vars_info = []
+        for i in range(Data.getVarCount()):
+            var_index = i # 0-based
+            name = Data.getVarName(var_index)
+            label = Data.getVarLabel(var_index)
+            type_str = Data.getVarType(var_index) # Returns int
+            
+            vars_info.append({
+                "name": name,
+                "label": label,
+                "type": str(type_str),
+            })
+        return vars_info
+
+    def get_variable_details(self, varname: str) -> str:
+        """Returns codebook/summary for a specific variable."""
+        return self.run_command(f"codebook {varname}")
+
+    def list_variables_structured(self) -> VariablesResponse:
+        vars_info: List[VariableInfo] = []
+        for item in self.list_variables():
+            vars_info.append(
+                VariableInfo(
+                    name=item.get("name", ""),
+                    label=item.get("label"),
+                    type=item.get("type"),
+                )
+            )
+        return VariablesResponse(variables=vars_info)
+
+    def list_graphs(self) -> List[str]:
+        """Returns list of graphs in memory."""
+        # 'graph dir' returns list in r(list)
+        # We need to ensure we run it quietly so we don't spam.
+        self.stata.run("quietly graph dir, memory")
+        
+        # Accessing r-class results in Python can be tricky via pystata's run command.
+        # We stash the result in a global macro that python sfi can easily read.
+        from sfi import Macro
+        self.stata.run("global mcp_graph_list `r(list)'")
+        graph_list_str = Macro.getGlobal("mcp_graph_list")
+        if not graph_list_str:
+            return []
+        
+        return graph_list_str.split()
+
+    def list_graphs_structured(self) -> GraphListResponse:
+        names = self.list_graphs()
+        active_name = names[-1] if names else None
+        graphs = [GraphInfo(name=n, active=(n == active_name)) for n in names]
+        return GraphListResponse(graphs=graphs)
+
+    def export_graph(self, graph_name: str = None, filename: str = None) -> str:
+        """Exports graph to a temp file and returns path."""
+        import tempfile
+        if not filename:
+            with tempfile.NamedTemporaryFile(prefix="mcp_stata_", suffix=".png", delete=False) as tmp:
+                filename = tmp.name
+        else:
+            # Ensure fresh start
+            if os.path.exists(filename):
+                try:
+                    os.remove(filename)
+                except Exception:
+                    pass
+            
+        cmd = "graph export"
+        if graph_name:
+            cmd += f' "{filename}", name("{graph_name}") replace'
+        else:
+            cmd += f' "{filename}", replace'
+            
+        output = self.run_command(cmd)
+        
+        if os.path.exists(filename):
+            try:
+                size = os.path.getsize(filename)
+                if size == 0:
+                    raise RuntimeError(f"Graph export failed: produced empty file {filename}")
+                if size > self.MAX_GRAPH_BYTES:
+                    raise RuntimeError(
+                        f"Graph export failed: file too large (> {self.MAX_GRAPH_BYTES} bytes): {filename}"
+                    )
+            except Exception as size_err:
+                # Clean up oversized or unreadable files
+                try:
+                    os.remove(filename)
+                except Exception:
+                    pass
+                raise size_err
+            return filename
+            
+        # If file missing, it failed. Check output for details.
+        raise RuntimeError(f"Graph export failed: {output}")
+
+    def get_help(self, topic: str, plain_text: bool = False) -> str:
+        """Returns help text as Markdown (default) or plain text."""
+        # Try to locate the .sthlp help file
+        # We use 'capture' to avoid crashing if not found
+        self.stata.run(f"capture findfile {topic}.sthlp")
+        
+        # Retrieve the found path from r(fn)
+        from sfi import Macro
+        self.stata.run("global mcp_help_file `r(fn)'")
+        fn = Macro.getGlobal("mcp_help_file")
+        
+        if fn and os.path.exists(fn):
+            try:
+                with open(fn, 'r', encoding='utf-8', errors='replace') as f:
+                    smcl = f.read()
+                if plain_text:
+                    return self._smcl_to_text(smcl)
+                try:
+                    return smcl_to_markdown(smcl, adopath=os.path.dirname(fn), current_file=os.path.splitext(os.path.basename(fn))[0])
+                except Exception as parse_err:
+                    logger.warning("SMCL to Markdown failed, falling back to plain text: %s", parse_err)
+                    return self._smcl_to_text(smcl)
+            except Exception as e:
+                return f"Error reading help file at {fn}: {e}"
+
+        # Fallback to URL if file not found
+        return f"Help file for '{topic}' not found. Please consult: https://www.stata.com/help.cgi?{topic}"
+
+    def get_stored_results(self) -> Dict[str, Any]:
+        """Returns e() and r() results."""
+        if not self._initialized:
+            self.init()
+            
+        from sfi import Scalar, Macro
+        
+        results = {"r": {}, "e": {}}
+        
+        # We parse 'return list' output as there is no direct bulk export of stored results
+        raw_r = self.run_command("return list")
+        raw_e = self.run_command("ereturn list")
+        
+        # Simple parser
+        def parse_list(text):
+            data = {}
+            # We don't strictly need to track sections if we check patterns
+            for line in text.splitlines():
+                line = line.strip()
+                if not line: continue
+                
+                # scalars: r(name) = value
+                if "=" in line and ("r(" in line or "e(" in line):
+                    try:
+                        name_part, val_part = line.split("=", 1)
+                        name_part = name_part.strip()  # "r(mean)"
+                        val_part = val_part.strip()    # "6165.2..."
+                        
+                        # Extract just the name inside r(...) if desired, 
+                        # or keep full key "r(mean)". 
+                        # User likely wants "mean" inside "r" dict.
+                        
+                        if "(" in name_part and name_part.endswith(")"):
+                            # r(mean) -> mean
+                            start = name_part.find("(") + 1
+                            end = name_part.find(")")
+                            key = name_part[start:end]
+                            data[key] = val_part
+                    except: pass
+                    
+                # macros: r(name) : "value"
+                elif ":" in line and ("r(" in line or "e(" in line):
+                     try:
+                        name_part, val_part = line.split(":", 1)
+                        name_part = name_part.strip()
+                        val_part = val_part.strip().strip('"')
+                        
+                        if "(" in name_part and name_part.endswith(")"):
+                            start = name_part.find("(") + 1
+                            end = name_part.find(")")
+                            key = name_part[start:end]
+                            data[key] = val_part
+                     except: pass
+            return data
+            
+        results["r"] = parse_list(raw_r)
+        results["e"] = parse_list(raw_e)
+        
+        return results
+
+    def export_graphs_all(self) -> GraphExportResponse:
+        """Exports all graphs to base64-encoded strings."""
+        exports: List[GraphExport] = []
+        for name in self.list_graphs():
+            try:
+                path = self.export_graph(name)
+                with open(path, "rb") as f:
+                    b64 = base64.b64encode(f.read()).decode("ascii")
+                exports.append(GraphExport(name=name, image_base64=b64))
+            except Exception as e:
+                logger.warning("Failed to export graph '%s': %s", name, e)
+                continue
+        return GraphExportResponse(graphs=exports)
+
+    def run_do_file(self, path: str, echo: bool = True, trace: bool = False) -> CommandResponse:
+        if not os.path.exists(path):
+            return CommandResponse(
+                command=f'do "{path}"',
+                rc=601,
+                stdout="",
+                stderr=None,
+                success=False,
+                error=ErrorEnvelope(
+                    message=f"Do-file not found: {path}",
+                    rc=601,
+                    command=path,
+                ),
+            )
+        return self._exec_with_capture(f'do "{path}"', echo=echo, trace=trace)
+
+    def load_data(self, source: str, clear: bool = True) -> CommandResponse:
+        src = source.strip()
+        clear_suffix = ", clear" if clear else ""
+
+        if src.startswith("sysuse "):
+            cmd = f"{src}{clear_suffix}"
+        elif src.startswith("webuse "):
+            cmd = f"{src}{clear_suffix}"
+        elif src.startswith("use "):
+            cmd = f"{src}{clear_suffix}"
+        elif "://" in src or src.endswith(".dta") or os.path.sep in src:
+            cmd = f'use "{src}"{clear_suffix}'
+        else:
+            cmd = f"sysuse {src}{clear_suffix}"
+
+        return self._exec_with_capture(cmd, echo=True, trace=False)
+
+    def codebook(self, varname: str, trace: bool = False) -> CommandResponse:
+        return self._exec_with_capture(f"codebook {varname}", trace=trace)
+