aigroup-stata-mcp 1.0.3__py3-none-any.whl

stata_mcp/core/stata/stata_finder/finder.py

@@ -0,0 +1,193 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import glob
+import logging
+import os
+import platform
+from typing import Optional
+
+from .linux import FinderLinux
+from .macos import FinderMacOS
+from .windows import FinderWindows, get_available_drives, windows_stata_match
+
+
+class StataFinder:
+   
+   
+   
+    FINDER_MAPPING = {
+        "Darwin": FinderMacOS,
+        "Windows": FinderWindows,
+        "Linux": FinderLinux,
+    }
+
+    def __init__(self, stata_cli: str = None):
+        """简要描述函数功能"""
+        finder_cls = self.FINDER_MAPPING.get(platform.system())
+        self.finder = finder_cls(stata_cli)
+
+    @property
+    def STATA_CLI(self) -> str | None:
+        """简要描述函数功能"""
+        try:
+            return self.finder.find_stata()
+        except (FileNotFoundError, AttributeError):
+            return None
+
+
+class StataFinderOLD:
+   
+   
+    """A class to find Stata CLI installations across different operating systems."""
+
+    FINDER_MAPPING = {
+        "Darwin": FinderMacOS,
+        "Windows": FinderWindows,
+        "Linux": FinderLinux,
+    }
+
+    def __init__(self, stata_cli: str = None):
+        """
+        Initialize the StataFinder.
+
+        Args:
+            stata_cli (str): the user updates stata_cli file path
+        """
+        self.current_os = platform.system()
+        self._os_finders = {
+            "Darwin": self._find_stata_macos,
+            "Windows": self._find_stata_windows,
+            "Linux": self._find_stata_linux,
+        }
+        self.finder = self._os_finders.get(self.current_os, None)
+        # TODO: Change the original finder to the newer
+        # self.finder = self.FINDER_MAPPING.get(self.current_os)(stata_cli = stata_cli)
+
+    # @property
+    # def STATA_CLI(self) -> str:  # 等前面的都改好了，就可以只保留这个了
+    #     return self.finder.find_stata()
+
+    def _stata_version_windows(self, driver: str = "C:\\") -> list:
+        """Find Stata installations on Windows."""
+        stata_paths = []
+        common_patterns = [
+            os.path.join(driver, "Program Files", "Stata*", "*.exe"),
+            os.path.join(driver, "Program Files(x86)", "Stata*", "*.exe"),
+        ]
+
+        for pattern in common_patterns:
+            try:
+                matches = glob.glob(pattern)
+                for match in matches:
+                    if "stata" in match.lower() and match.lower().endswith(".exe"):
+                        stata_paths.append(match)
+
+                if not stata_paths:
+                    for root, dirs, files in os.walk(driver):
+                        if root.count(os.sep) - driver.count(os.sep) > 3:
+                            dirs.clear()
+                            continue
+
+                        for file in files:
+                            if (
+                                file.lower().endswith(".exe")
+                                and "stata" in file.lower()
+                            ):
+                                stata_paths.append(os.path.join(root, file))
+
+            except Exception as e:
+                logging.warn(e)
+
+        return stata_paths
+
+    def _find_stata_macos(self) -> Optional[str]:
+        return self.finder.find_stata()
+
+    def _default_stata_cli_path_windows(self) -> Optional[str]:
+        """Get default Stata CLI path on Windows."""
+        drives = get_available_drives()
+        stata_cli_path_list = []
+
+        for drive in drives:
+            stata_cli_path_list += self._stata_version_windows(drive)
+
+        if len(stata_cli_path_list) == 0:
+            return None
+        elif len(stata_cli_path_list) == 1:
+            return stata_cli_path_list[0]
+        else:
+            for path in stata_cli_path_list:
+                if windows_stata_match(path):
+                    return path
+            return stata_cli_path_list[0]
+
+    def _find_stata_windows(self) -> Optional[str]:
+        """Find Stata CLI on Windows systems."""
+        return self._default_stata_cli_path_windows()
+
+    def _default_stata_cli_path_linux(self) -> Optional[str]:
+        """Get default Stata CLI path on Linux."""
+        # TODO: Implement Linux-specific logic
+        return None
+
+    def _find_stata_linux(self) -> Optional[str]:
+        """Find the Stata CLI path on Linux systems.
+
+        For Linux users, this function attempts to locate the Stata CLI executable.
+
+        Returns:
+            The path to the Stata CLI executable, or None if not found.
+        """
+        return self._default_stata_cli_path_linux()
+
+    def find_stata(self,
+                   os_name: Optional[str] = None,
+                   is_env: bool = True) -> Optional[str]:
+        """Find Stata CLI installation.
+
+        Args:
+            os_name: Operating system name. If None, uses current system.
+            is_env: Whether to check environment variables first.
+
+        Returns:
+            Path to Stata CLI executable, or None if not found.
+
+        Raises:
+            RuntimeError: If the operating system is not supported.
+        """
+        if is_env:
+            stata_cli = os.getenv("stata_cli", None) or os.getenv("STATA_CLI", None)
+            if stata_cli:
+                return stata_cli
+
+        target_os = os_name or self.current_os
+        finder = self._os_finders.get(target_os)
+
+        if not finder:
+            raise RuntimeError(f"Unsupported OS: {target_os!r}")
+
+        return finder()
+
+    def get_supported_os(self) -> list:
+        """Get list of supported operating systems."""
+        return list(self._os_finders.keys())
+
+    def is_stata_available(
+        self, os_name: Optional[str] = None, is_env: bool = True
+    ) -> bool:
+        """Check if Stata is available on the system.
+
+        Args:
+            os_name: Operating system name. If None, uses current system.
+            is_env: Whether to check environment variables first.
+
+        Returns:
+            True if Stata is found, False otherwise.
+        """
+        try:
+            stata_path = self.find_stata(os_name=os_name, is_env=is_env)
+            return stata_path is not None
+        except RuntimeError:
+            return False

stata_mcp/core/stata/stata_finder/linux.py

@@ -0,0 +1,43 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+from pathlib import Path
+from typing import Dict, List
+
+from .base import FinderBase
+
+
+class FinderLinux(FinderBase):
+   
+   
+   
+    def finder(self) -> str:
+        """简要描述函数功能"""
+        bin_results = self.find_from_bin()
+        if bin_results:
+            return max(bin_results).stata_cli_path
+        else:
+            raise FileNotFoundError("Stata CLI not found")
+
+    def find_path_base(self) -> Dict[str, List[str]]:
+        """简要描述函数功能"""
+        # Start with default bin directory
+        bin_dirs = ["/usr/local/bin"]
+
+        # Search for additional directories containing "stata" in /usr/local/bin
+        usr_local_bin = Path("/usr/local/bin")
+        if usr_local_bin.exists() and usr_local_bin.is_dir():
+            # Look for directories containing "stata" in their name
+            for item in usr_local_bin.iterdir():
+                if item.is_dir() and "stata" in item.name.lower():
+                    # Add the stata directory path to search directories
+                    bin_dirs.append(str(item))
+
+        return {
+            "bin": bin_dirs,
+        }
+
+
+if __name__ == "__main__":
+    finder = FinderLinux()

stata_mcp/core/stata/stata_finder/macos.py

@@ -0,0 +1,88 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import re
+from pathlib import Path
+from typing import Dict, List
+
+from .base import FinderBase, StataEditionConfig
+
+
+class FinderMacOS(FinderBase):
+   
+   
+   
+    def finder(self) -> str | None:
+        """简要描述函数功能"""
+        bin_results = self.find_from_bin()
+        if bin_results:
+            return max(bin_results).stata_cli_path
+
+        application_results = self.find_from_application()
+        if application_results:
+            return max(application_results).stata_cli_path
+        else:  # If there is no Stata CLI found, raise an error
+            raise FileNotFoundError("Stata CLI not found")
+
+    def find_path_base(self) -> Dict[str, List[str]]:
+        """简要描述函数功能"""
+        return {
+            "bin": ["/usr/local/bin"],
+            "application": ["/Applications"],
+        }
+
+    def _application_find_base(self,
+                               dot_app: str | Path,
+                               version: int | float = None) -> StataEditionConfig | None:
+        _version = version
+        _edition = None
+        stata_cli_path = None
+
+        if not _version:
+            for isstata_file in dot_app.glob("isstata.*"):
+                if isstata_file.is_file():
+                    # Extract version number from filename like "isstata.180"
+                    match = re.search(r'isstata\.(\d+)', isstata_file.name.lower())
+                    if match:
+                        _version = float(match.group(1)) / 10
+                        break
+        for stata_app in dot_app.glob("Stata*.app"):
+            if stata_app.is_dir():
+                # Extract edition from Stata app name (MP, SE, BE, IC)
+                # Remove "Stata" prefix and ".app" suffix, then convert to lowercase
+                _edition = stata_app.name.replace("Stata", "").replace(".app", "").lower()
+                __stata_cli_path = stata_app / "Contents" / "MacOS" / f"stata-{_edition}"
+                if self._is_executable(__stata_cli_path):
+                    stata_cli_path = str(__stata_cli_path)
+                    break
+        if _version and _edition and stata_cli_path:
+            return StataEditionConfig(_edition, _version, stata_cli_path)
+
+        else:
+            return None
+
+    def find_from_application(self) -> List[StataEditionConfig]:
+        """简要描述函数功能"""
+        found_executables: List[StataEditionConfig] = []
+        applications_dir = Path(self.find_path_base().get("application")[0])
+
+        # Check for /Applications/Stata directory for Multi-Stata Exist
+        stata_dir = applications_dir / "Stata"
+        if default_stata := self._application_find_base(stata_dir):  # If exist default, return directly.
+            return [default_stata]
+
+        # 通过for循环来从applications_dir里找stata*.app
+        for stata_app in applications_dir.glob("Stata *"):
+            _version = None
+            if stata_app.is_dir():
+                _version = eval(stata_app.name.split()[-1])
+                if stata_app_config := self._application_find_base(stata_app, version=_version):
+                    found_executables.append(stata_app_config)
+
+        return found_executables
+
+
+if __name__ == "__main__":
+    finder = FinderMacOS()
+    print(finder.finder())

stata_mcp/core/stata/stata_finder/windows.py

@@ -0,0 +1,191 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import glob
+import os
+import re
+import string
+from typing import Dict, List
+
+from .base import FinderBase, StataEditionConfig
+
+
+def get_available_drives():
+    """简要描述函数功能"""
+    drives = []
+    for letter in string.ascii_uppercase:
+        if os.path.exists(f"{letter}:\\"):
+            drives.append(f"{letter}:\\")
+    return drives
+
+
+def windows_stata_match(path: str) -> bool:
+    """
+    Check whether the given path matches the pattern of a Windows
+    Stata executable.
+
+    Args:
+        path: Path string to be checked.
+
+    Returns:
+        bool: ``True`` if the path matches a Stata executable pattern,
+        otherwise ``False``.
+    """
+    # Regular expression matching ``Stata\d+\Stata(MP|SE|BE|IC)?.exe``
+    # ``\d+`` matches one or more digits (the version number)
+    # ``(MP|SE|BE|IC)?`` matches an optional edition suffix
+    pattern = r"Stata\d+\\\\Stata(MP|SE|BE|IC)?\.exe$"
+
+    if re.search(pattern, path):
+        return True
+    return False
+
+
+class FinderWindows(FinderBase):
+   
+   
+   
+    def finder(self) -> str:
+        """简要描述函数功能"""
+        default_results = self.find_from_default_install_path()
+        if default_results:
+            return max(default_results).stata_cli_path
+
+        driver_results = self.scan_stata_from_drivers()
+        if driver_results:
+            return max(driver_results).stata_cli_path
+
+        deep_results = self.scan_stata_deeply()
+        if deep_results:
+            return max(deep_results).stata_cli_path
+        else:
+            raise FileNotFoundError("Stata executable not found")
+
+    def find_path_base(self) -> Dict[str, List[str]]:
+        """简要描述函数功能"""
+        return {
+            "default": [
+                r"C:\Program Files\Stata*",
+                r"C:\Program Files (x86)\Stata*",
+                r"D:\Program Files\Stata*",
+                r"D:\Program Files (x86)\Stata*",
+            ],
+            "drivers": get_available_drives()
+        }
+
+    def find_from_default_install_path(self) -> List[StataEditionConfig]:
+        """Find Stata installations in default Windows installation paths.
+
+        This function searches for Stata in standard Windows locations
+        like Program Files directories using simple matching.
+
+        Returns:
+            List of StataEditionConfig objects found in default paths
+        """
+        found_configs = []
+
+        # Common Windows installation paths
+        common_paths = self.find_path_base().get("default")
+
+        for path_pattern in common_paths:
+            try:
+                matches = glob.glob(path_pattern)
+                for match in matches:
+                    # Search in these Stata directories for executables
+                    executables = glob.glob(os.path.join(match, "*.exe"))
+                    for exe in executables:
+                        # Simple matching like the original implementation
+                        if "stata" in exe.lower() and exe.lower().endswith(".exe"):
+                            config = StataEditionConfig.from_path(exe)
+                            found_configs.append(config)
+            except Exception:
+                pass
+
+        return found_configs
+
+    def scan_stata_from_drivers(self) -> List[StataEditionConfig]:
+        """Scan all available drivers for Stata installations in non-standard locations.
+
+        This method searches for Stata in locations that are NOT in standard
+        installation paths (Program Files, etc.), focusing on custom or
+        alternative installation locations.
+
+        Returns:
+            List of StataEditionConfig objects found in non-standard locations
+        """
+        drivers_base = self.find_path_base().get("drivers")
+        found_configs = []
+
+        if not drivers_base:
+            return found_configs
+
+        for driver in drivers_base:
+            # Skip standard installation directories to avoid duplication
+            skip_patterns = [
+                os.path.join(driver, "Program Files"),
+                os.path.join(driver, "Program Files (x86)"),
+            ]
+
+            try:
+                # Limit search depth to avoid excessive scanning
+                for root, dirs, files in os.walk(driver):
+                    # Skip standard installation directories
+                    should_skip = False
+                    for skip_pattern in skip_patterns:
+                        if root.lower().startswith(skip_pattern.lower()):
+                            should_skip = True
+                            break
+
+                    if should_skip:
+                        continue
+
+                    # Limit depth to 4 levels
+                    if root.count(os.sep) - driver.count(os.sep) > 4:
+                        dirs.clear()
+                        continue
+
+                    for file in files:
+                        if (
+                            file.lower().endswith(".exe")
+                            and "stata" in file.lower()
+                        ):
+                            full_path = os.path.join(root, file)
+                            config = StataEditionConfig.from_path(full_path)
+                            found_configs.append(config)
+            except Exception:
+                pass
+
+        return found_configs
+
+    def scan_stata_deeply(self) -> List[StataEditionConfig]:
+        """Perform deep scan across all drives for any Stata installation.
+
+        This is the final fallback search method that scans everywhere
+        without restrictions, including standard paths, to find any possible
+        Stata installation that might have been missed.
+
+        Returns:
+            List of StataEditionConfig objects found through deep scanning
+        """
+        drivers_base = self.find_path_base().get("drivers")
+        found_configs = []
+
+        if not drivers_base:
+            return found_configs
+
+        for driver in drivers_base:
+            try:
+                # Deep scan without restrictions
+                for root, dirs, files in os.walk(driver):
+                    for file in files:
+                        if file.lower().endswith(".exe") and "stata" in file.lower():
+                            full_path = os.path.join(root, file)
+                            # Less strict matching for deep scan
+                            if "stata" in full_path.lower():
+                                config = StataEditionConfig.from_path(full_path)
+                                found_configs.append(config)
+            except Exception:
+                pass
+
+        return found_configs

stata_mcp/server/__init__.py

@@ -0,0 +1,8 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+"""Stata MCP Server package."""
+from .main import create_stata_server
+
+__all__ = ["create_stata_server"]

stata_mcp/server/main.py

@@ -0,0 +1,153 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+import locale
+import os
+import platform
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from mcp.server.fastmcp import Context, FastMCP, Icon
+from mcp.server.session import ServerSession
+from pydantic import BaseModel, Field
+
+from ..core.data_info import CsvDataInfo, DtaDataInfo
+from ..core.stata import StataController, StataDo, StataFinder
+from ..utils.Prompt import pmp
+
+from .tools.core_tools import register_core_tools
+from .tools.file_tools import register_file_tools
+from .tools.stata_tools import register_stata_tools
+from .prompts.core_prompts import register_core_prompts
+
+
+class StataServerConfig(BaseModel):
+   
+   
+    """Configuration for Stata MCP server."""
+    
+    name: str = "stata-mcp"
+    instructions: str = "Stata-MCP lets you and LLMs run Stata do-files and fetch results"
+    website_url: str = "https://github.com/jackdark425"
+    working_directory: Optional[str] = None
+    stata_cli: Optional[str] = None
+
+
+class StataServerContext(BaseModel):
+   
+   
+    """Application context with typed dependencies."""
+    
+    config: StataServerConfig
+    stata_finder: Any = Field(description="StataFinder instance")
+    working_directory: Path
+    output_base_path: Path
+
+    class Config:
+        arbitrary_types_allowed = True
+
+
+@asynccontextmanager
+async def server_lifespan(server: FastMCP, config: StataServerConfig) -> AsyncIterator[StataServerContext]:
+    """Manage server startup and shutdown lifecycle."""
+    # Initialize system configuration
+    system_os = platform.system()
+    if system_os not in ["Darwin", "Linux", "Windows"]:
+        raise RuntimeError(f"Unsupported operating system: {system_os}")
+    
+    # Find Stata CLI
+    stata_finder = StataFinder()
+    stata_cli = config.stata_cli or stata_finder.STATA_CLI
+    
+    # Determine working directory
+    if config.working_directory:
+        working_directory = Path(config.working_directory)
+    else:
+        client = os.getenv("STATA-MCP-CLIENT")
+        if client == "cc":  # Claude Code
+            working_directory = Path.cwd()
+        else:
+            cwd = os.getenv("STATA_MCP_CWD", os.getenv("STATA-MCP-CWD", None))
+            if cwd:
+                working_directory = Path(cwd)
+            else:
+                if system_os in ["Darwin", "Linux"]:
+                    working_directory = Path.home() / "Documents"
+                else:
+                    working_directory = Path(os.getenv("USERPROFILE", "~")) / "Documents"
+    
+    # Create output directories
+    output_base_path = working_directory / "stata-mcp-folder"
+    output_base_path.mkdir(exist_ok=True)
+    
+    # Create subdirectories
+    log_base_path = output_base_path / "stata-mcp-log"
+    dofile_base_path = output_base_path / "stata-mcp-dofile"
+    result_doc_path = output_base_path / "stata-mcp-result"
+    tmp_base_path = output_base_path / "stata-mcp-tmp"
+    
+    for path in [log_base_path, dofile_base_path, result_doc_path, tmp_base_path]:
+        path.mkdir(exist_ok=True)
+    
+    # Set language for prompts
+    lang_mapping = {"zh-CN": "cn", "en_US": "en"}
+    lang, _ = locale.getdefaultlocale()
+    pmp.set_lang(lang_mapping.get(lang, "en"))
+    
+    yield StataServerContext(
+        config=config,
+        stata_finder=stata_finder,
+        working_directory=working_directory,
+        output_base_path=output_base_path
+    )
+
+
+def create_stata_server(config: Optional[StataServerConfig] = None) -> FastMCP:
+    """Create and configure the Stata MCP server."""
+    config = config or StataServerConfig()
+    
+    # Create server with lifespan
+    @asynccontextmanager
+    async def lifespan(server: FastMCP) -> AsyncIterator[Dict[str, Any]]:
+        async with server_lifespan(server, config) as context:
+            yield {"stata_context": context}
+    
+    # Initialize server with latest MCP features
+    server = FastMCP(
+        name=config.name,
+        instructions=config.instructions,
+        website_url=config.website_url,
+        icons=[
+            Icon(
+                src="https://avatars.githubusercontent.com/u/201514154?v=4",
+                mimeType="image/png",
+                sizes=["460*460"]
+            )
+        ],
+        lifespan=lifespan
+    )
+    
+    # Register all components
+    register_core_tools(server)
+    register_file_tools(server)
+    register_stata_tools(server)
+    register_core_prompts(server)
+    
+    return server
+
+
+# Export the default server instance
+stata_server = create_stata_server()
+
+
+def run_server(transport: str = "stdio") -> None:
+    """Run the Stata MCP server with the specified transport."""
+    stata_server.run(transport=transport)
+
+
+if __name__ == "__main__":
+    run_server()

stata_mcp/server/prompts/__init__.py

@@ -0,0 +1,8 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+"""Stata MCP Prompts package."""
+from .core_prompts import register_core_prompts
+
+__all__ = ["register_core_prompts"]