simstack 2.0.1__py3-none-any.whl

simstack/util/db_logger.py

@@ -0,0 +1,117 @@
+import logging
+import re
+from datetime import datetime
+from typing import Optional, Dict, Any
+
+from pymongo import MongoClient
+
+# Regular expression patterns
+task_pattern = re.compile(r"task_id: ([\w\-\.]+)")
+resource_pattern = re.compile(r"resource_id: ([\w\-\.]+)")
+
+
+def extract_task_id(message: str) -> Optional[str]:
+    """Extract task ID from log message."""
+    match = task_pattern.search(message)
+    return match.group(1) if match else None
+
+
+def extract_resource_id(message: str) -> Optional[str]:
+    """Extract resource ID from log message."""
+    match = resource_pattern.search(message)
+    return match.group(1) if match else None
+
+
+class DBLogHandler(logging.Handler):
+    """Logging handler that stores logs in MongoDB using pymongo."""
+
+    def __init__(
+        self,
+        connection_string: str,
+        db_name: str,
+        collection_name: str = "logs",
+    ):
+        super().__init__()
+
+        # Validate configuration
+        if not connection_string or not isinstance(connection_string, str):
+            raise ValueError(f"Invalid connection string {connection_string}")
+        if not db_name or not isinstance(db_name, str):
+            raise ValueError(f"Invalid database name {db_name}")
+        if not collection_name or not isinstance(collection_name, str):
+            raise ValueError(f"Invalid collection name {collection_name}")
+
+        # Store configuration
+        self.connection_string = connection_string
+        self.db_name = db_name
+        self.collection_name = collection_name
+
+        # Create MongoDB client
+        self.client = MongoClient(self.connection_string)
+        self.db = self.client[self.db_name]
+        self.collection = self.db[self.collection_name]
+
+    def emit(self, record: logging.LogRecord) -> None:
+        """Process a log record by converting it to a log entry and inserting it into MongoDB."""
+        try:
+            # Create log entry as a dictionary
+            log_entry = self.create_log_entry(record)
+
+            # Insert the log entry into the collection
+            self.collection.insert_one(log_entry)
+
+        except Exception as e:
+            # Use print instead of logging to avoid potential recursion
+            print(f"Error in emit: {e}")
+            self.handleError(record)
+
+    def create_log_entry(self, record: logging.LogRecord) -> Dict[str, Any]:
+        # Format message
+        message = self.format(record)
+
+        # Extract task and resource IDs
+        task_id = extract_task_id(message)
+        resource_id = extract_resource_id(message)
+
+        # Create log entry dictionary
+        log_entry_dict = {
+            "timestamp": datetime.fromtimestamp(record.created),
+            "level": record.levelname,
+            "logger_name": record.name,
+            "message": message,
+            "module": record.module,
+            "function": record.funcName,
+            "line": record.lineno,
+            "task_id": task_id,
+            "resource": resource_id,
+            "thread_name": record.threadName,
+            "process_name": record.processName,
+        }
+
+        # Add exception info if available
+        if record.exc_info:
+            log_entry_dict.update(
+                {
+                    "exception_type": record.exc_info[0].__name__
+                    if record.exc_info and record.exc_info[0]
+                    else None,
+                    "exception_message": str(record.exc_info[1])
+                    if record.exc_info and record.exc_info[1]
+                    else None,
+                    "exception_traceback": record.exc_text,
+                }
+            )
+
+        return log_entry_dict
+
+    def close(self) -> None:
+        """Close the MongoDB client connection."""
+        try:
+            if hasattr(self, "client") and self.client:
+                self.client.close()
+        except Exception as e:
+            print(f"Error closing DB log handler client: {e}")
+
+    def __del__(self):
+        """Ensure connection is closed when handler is destroyed."""
+        self.close()

simstack/util/default_from_dict.py

@@ -0,0 +1,57 @@
+import logging
+from typing import Any
+
+from odmantic import Model
+
+logger = logging.getLogger(__name__)
+
+
+def default_from_model(cls_method, model: Model, **kwargs) -> Model:
+    model_dict = model.model_dump(exclude={"id"})
+    return cls_method.from_dict(model_dict, **kwargs)
+
+
+def default_from_dict(cls_method, data: dict, **kwargs) -> Any:
+    """
+    Create an instance of the model from a dictionary.
+    Handles nested models and enum values.
+    """
+    if not isinstance(data, dict):
+        return data  # Return as is if it's not a dictionary
+
+    processed_data = {}
+    for field_name, field in cls_method.model_fields.items():
+        if field_name not in data:
+            continue
+
+        field_type = field.annotation
+        field_value = data[field_name]
+
+        # Handle enum values
+        if hasattr(field_type, "__members__"):
+            logger.debug("Handling enum field: %s", field_name)
+            if isinstance(field_value, field_type):
+                # If it's already an enum instance, use it directly
+                logger.debug("Field value is enum instance: %s", field_value)
+                processed_data[field_name] = field_value
+            else:
+                # Try to create enum from the value
+                try:
+                    processed_data[field_name] = field_type(field_value)
+                    logger.debug(
+                        "Converted field value to enum: %s", processed_data[field_name]
+                    )
+                except ValueError:
+                    logger.debug(
+                        "Invalid enum value for field %s: %s", field_name, field_value
+                    )
+                    raise ValueError(
+                        f"Invalid enum value '{field_value}' for field {field_name}"
+                    )
+        # Handle nested models
+        elif hasattr(field_type, "from_dict") and isinstance(field_value, dict):
+            processed_data[field_name] = field_type.from_dict(field_value)
+        else:
+            processed_data[field_name] = field_value
+
+    return cls_method.model_validate(processed_data)

simstack/util/directory_iterator.py

@@ -0,0 +1,216 @@
+import fnmatch
+import os
+from pathlib import Path
+from typing import Iterator, List, Union, Optional
+
+
+class DirectoryPath:
+    """
+    Represents a directory path with excluded subdirectory patterns and provides iteration capabilities.
+    """
+
+    def __init__(
+        self,
+        root_path: Union[str, Path],
+        excluded_patterns: Optional[List[str]] = None,
+        file_extensions: Optional[List[str]] = None,
+    ):
+        """
+        Initialize the DirectoryPath.
+
+        Args:
+            root_path: The root directory path to traverse
+            excluded_patterns: List of patterns to exclude (supports wildcards like *, ?, [])
+            file_extensions: List of file extensions to include (e.g., ['.py', '.txt'])
+        """
+        self.root_path = Path(root_path).resolve()
+        self.excluded_patterns = excluded_patterns or []
+        self.file_extensions = file_extensions or []
+
+        if not self.root_path.exists():
+            raise FileNotFoundError(f"Directory does not exist: {self.root_path}")
+        if not self.root_path.is_dir():
+            raise NotADirectoryError(f"Path is not a directory: {self.root_path}")
+
+    def _is_excluded(self, path: Path) -> bool:
+        """
+        Check if a path matches any of the excluded patterns.
+
+        Args:
+            path: Path to check
+
+        Returns:
+            True if the path should be excluded, False otherwise
+        """
+        path_str = str(path)
+        relative_path = str(path.relative_to(self.root_path))
+
+        for pattern in self.excluded_patterns:
+            # Check against full path
+            if fnmatch.fnmatch(path_str, pattern):
+                return True
+            # Check against relative path
+            if fnmatch.fnmatch(relative_path, pattern):
+                return True
+            # Check against directory name only
+            if fnmatch.fnmatch(path.name, pattern):
+                return True
+
+        return False
+
+    def _matches_extensions(self, file_path: Path) -> bool:
+        """
+        Check if a file matches the specified extensions.
+
+        Args:
+            file_path: File path to check
+
+        Returns:
+            True if the file matches the extensions or no extensions specified
+        """
+        if not self.file_extensions:
+            return True
+
+        return file_path.suffix.lower() in [ext.lower() for ext in self.file_extensions]
+
+    def iterate_files(self) -> Iterator[Path]:
+        """
+        Iterate over all files in the directory tree, respecting exclusion patterns.
+
+        Yields:
+            Path objects for each file found
+        """
+        for root, dirs, files in os.walk(self.root_path):
+            root_path = Path(root)
+
+            # Filter out excluded directories in-place to prevent walking into them
+            dirs[:] = [d for d in dirs if not self._is_excluded(root_path / d)]
+
+            # Yield files that match criteria
+            for file in files:
+                file_path = root_path / file
+                if not self._is_excluded(file_path) and self._matches_extensions(
+                    file_path
+                ):
+                    yield file_path
+
+    def iterate_directories(self) -> Iterator[Path]:
+        """
+        Iterate over all directories in the directory tree, respecting exclusion patterns.
+
+        Yields:
+            Path objects for each directory found
+        """
+        for root, dirs, _ in os.walk(self.root_path):
+            root_path = Path(root)
+
+            # Filter and yield directories
+            for dir_name in dirs:
+                dir_path = root_path / dir_name
+                if not self._is_excluded(dir_path):
+                    yield dir_path
+
+            # Filter out excluded directories to prevent walking into them
+            dirs[:] = [d for d in dirs if not self._is_excluded(root_path / d)]
+
+    def __iter__(self) -> Iterator[Path]:
+        """
+        Default iterator returns files.
+        """
+        return self.iterate_files()
+
+    def get_files_list(self) -> List[Path]:
+        """
+        Get a list of all files matching the criteria.
+
+        Returns:
+            List of Path objects
+        """
+        return list(self.iterate_files())
+
+    def get_directories_list(self) -> List[Path]:
+        """
+        Get a list of all directories matching the criteria.
+
+        Returns:
+            List of Path objects
+        """
+        return list(self.iterate_directories())
+
+    def add_excluded_pattern(self, pattern: str) -> None:
+        """
+        Add an exclusion pattern.
+
+        Args:
+            pattern: Pattern to exclude
+        """
+        if pattern not in self.excluded_patterns:
+            self.excluded_patterns.append(pattern)
+
+    def remove_excluded_pattern(self, pattern: str) -> None:
+        """
+        Remove an exclusion pattern.
+
+        Args:
+            pattern: Pattern to remove
+        """
+        if pattern in self.excluded_patterns:
+            self.excluded_patterns.remove(pattern)
+
+    def __str__(self) -> str:
+        return (
+            f"DirectoryPath(root='{self.root_path}', excluded={self.excluded_patterns})"
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"DirectoryPath(root_path='{self.root_path}', "
+            f"excluded_patterns={self.excluded_patterns}, "
+            f"file_extensions={self.file_extensions})"
+        )
+
+
+# Example usage and utility functions
+def find_python_files_with_exclusions(
+    root_path: str, excluded_patterns: List[str] = None
+) -> List[Path]:
+    """
+    Find Python files using the DirectoryPath class with exclusions.
+
+    Args:
+        root_path: Root directory to search
+        excluded_patterns: Patterns to exclude
+
+    Returns:
+        List of Python file paths
+    """
+    excluded_patterns = excluded_patterns or [
+        "__pycache__",
+        "*.pyc",
+        ".git",
+        ".venv",
+        "venv",
+    ]
+    dir_path = DirectoryPath(root_path, excluded_patterns, [".py"])
+    return dir_path.get_files_list()
+
+
+if __name__ == "__main__":
+    # Example usage
+    try:
+        # Create a DirectoryPath instance
+        dp = DirectoryPath(
+            root_path=".",
+            excluded_patterns=["__pycache__", "*.pyc", ".git", "node_modules", "*.log"],
+            file_extensions=[".py", ".txt"],
+        )
+
+        print(f"Directory Path: {dp}")
+        print("\nFound files:")
+        for file_path in dp:
+            print(f"  {file_path}")
+
+        print(f"\nTotal files found: {len(dp.get_files_list())}")
+
+    except (FileNotFoundError, NotADirectoryError) as e:
+        print(f"Error: {e}")

simstack/util/docstring_parser.py

@@ -0,0 +1,186 @@
+import inspect
+import re
+from typing import Optional, Dict, Any, List
+
+
+class DocstringParser:
+    """
+    Parse a docstring into structured components.
+
+    All component getters return None when the corresponding section is not present.
+    Supported sections:
+      - Args: / Parameters:
+      - Returns:
+      - SimstackResult:
+      - CalledNodes:
+      - Raises:
+    """
+
+    _SECTION_NAMES = ("Args:", "Parameters:", "Returns:", "SimstackResult:", "CalledNodes:", "Raises:")
+
+    def __init__(self, docstring: Optional[str]) -> None:
+        self._raw = docstring or ""
+        self._doc = inspect.cleandoc(self._raw) if self._raw else ""
+        self._cache: Dict[str, Any] = {}
+
+    def _section(self, name: str, until: tuple[str, ...]) -> Optional[str]:
+        """
+        Extract the body text after `name` up to any of the `until` markers (or end).
+        Returns None if the section marker isn't present.
+        """
+        if not self._doc:
+            return None
+
+        # e.g. r"Returns:(.*?)(?:SimstackResult:|CalledNodes:|Raises:|$)"
+        until_alt = "|".join(map(re.escape, until)) + "|$"
+        pattern = rf"{re.escape(name)}(.*?)(?:{until_alt})"
+        match = re.search(pattern, self._doc, re.DOTALL)
+        if not match:
+            return None
+
+        body = match.group(1).strip()
+        return body if body else None
+
+    def description(self) -> Optional[str]:
+        if "description" in self._cache:
+            return self._cache["description"]
+
+        if not self._doc:
+            self._cache["description"] = None
+            return None
+
+        match = re.search(
+            r"^(.*?)(?:Args:|Parameters:|Returns:|SimstackResult:|CalledNodes:|Raises:|$)",
+            self._doc,
+            re.DOTALL,
+        )
+        desc = match.group(1).strip() if match else ""
+        self._cache["description"] = desc if desc else None
+        return self._cache["description"]
+
+    def params(self) -> Optional[Dict[str, Dict[str, Optional[str]]]]:
+        if "params" in self._cache:
+            return self._cache["params"]
+
+        param_text = self._section("Args:", ("Returns:", "SimstackResult:", "CalledNodes:", "Raises:"))
+        if param_text is None:
+            param_text = self._section("Parameters:", ("Returns:", "SimstackResult:", "CalledNodes:", "Raises:"))
+
+        if not param_text:
+            self._cache["params"] = None
+            return None
+
+        params: Dict[str, Dict[str, Optional[str]]] = {}
+        param_matches = re.finditer(
+            r"(\w+)\s*(?:\(([^)]+)\))?\s*:\s*(.+?)(?=\n\s*\w+\s*:|$)",
+            param_text,
+            re.DOTALL,
+        )
+        for match in param_matches:
+            param_name = match.group(1)
+            param_type = match.group(2)  # Maybe None
+            param_desc = match.group(3).strip()
+            params[param_name] = {"type": param_type, "description": param_desc}
+
+        self._cache["params"] = params if params else None
+        return self._cache["params"]
+
+    def returns(self) -> Optional[Dict[str, str]]:
+        if "returns" in self._cache:
+            return self._cache["returns"]
+
+        return_text = self._section("Returns:", ("SimstackResult:", "CalledNodes:", "Raises:"))
+        if not return_text:
+            self._cache["returns"] = None
+            return None
+
+        self._cache["returns"] = {"description": return_text}
+        return self._cache["returns"]
+
+    def simstack_results(self) -> Optional[Dict[str, Dict[str, str]]]:
+        if "simstack_results" in self._cache:
+            return self._cache["simstack_results"]
+
+        simstack_text = self._section("SimstackResult:", ("CalledNodes:", "Raises:"))
+        if not simstack_text:
+            self._cache["simstack_results"] = None
+            return None
+
+        simstack_results: Dict[str, Dict[str, str]] = {}
+        simstack_matches = re.finditer(
+            r"(\w+)\s*\(([^)]+)\)\s*(.+?)(?=\n\s*\w+\s*\(|$)",
+            simstack_text,
+            re.DOTALL,
+        )
+        for match in simstack_matches:
+            result_name = match.group(1)
+            result_type = match.group(2).strip()
+            result_desc = match.group(3).strip()
+            simstack_results[result_name] = {"name": result_name, "type": result_type, "description": result_desc}
+
+        self._cache["simstack_results"] = simstack_results if simstack_results else None
+        return self._cache["simstack_results"]
+
+    def called_nodes(self) -> Optional[List[str]]:
+        if "called_nodes" in self._cache:
+            return self._cache["called_nodes"]
+
+        called_nodes_text = self._section("Called Nodes:", ("Raises:","Notes"))
+        if not called_nodes_text:
+            self._cache["called_nodes"] = None
+            return None
+
+        nodes = [
+            line.strip().lstrip("-").strip()
+            for line in called_nodes_text.split("\n")
+            if line.strip()
+        ]
+        self._cache["called_nodes"] = nodes if nodes else None
+        return self._cache["called_nodes"]
+
+    def raises(self) -> Optional[Dict[str, Dict[str, str]]]:
+        if "raises" in self._cache:
+            return self._cache["raises"]
+
+        raises_text = self._section("Raises:", ())
+        if not raises_text:
+            self._cache["raises"] = None
+            return None
+
+        raises: Dict[str, Dict[str, str]] = {}
+        raises_matches = re.finditer(
+            r"(\w+)\s*:\s*(.+?)(?=\n\s*\w+\s*:|$)",
+            raises_text,
+            re.DOTALL,
+        )
+        for match in raises_matches:
+            exception_name = match.group(1)
+            exception_desc = match.group(2).strip()
+            raises[exception_name] = {"description": exception_desc}
+
+        self._cache["raises"] = raises if raises else None
+        return self._cache["raises"]
+
+    def as_dict(self) -> Dict[str, Any]:
+        """
+        Backward-compatible structure (always includes keys),
+        while still using the None-returning component getters internally.
+        """
+        return {
+            "description": self.description() or "",
+            "params": self.params() or {},
+            "returns": self.returns() or {},
+            "simstack_results": self.simstack_results() or {},
+            "called_nodes": self.called_nodes() or [],
+            "raises": self.raises() or {},
+        }
+
+
+def parse_docstring(docstring: Optional[str]) -> Dict[str, Any]:
+    """
+    Backward-compatible wrapper around DocstringParser.
+
+    Prefer using DocstringParser directly when you want per-section access with
+    None when a section is missing.
+    """
+    return DocstringParser(docstring).as_dict()

simstack/util/file_hashing.py

@@ -0,0 +1,72 @@
+import hashlib
+from pathlib import Path
+from typing import Union, BinaryIO
+
+def hash_file(
+    file_path: Union[str, Path], algorithm: str = "sha256", chunk_size: int = 8192
+) -> str:
+    """
+    Calculate hash of a file on disk.
+
+    Args:
+        file_path: Path to the file
+        algorithm: Hash algorithm (md5, sha1, sha256, etc.)
+        chunk_size: Size of chunks to read
+
+    Returns:
+        str: Hexadecimal digest of the hash
+    """
+    hash_obj = hashlib.new(algorithm)
+
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            hash_obj.update(chunk)
+
+    return hash_obj.hexdigest()
+
+def hash_string(string: str, algorithm: str = "sha256") -> str:
+    hash_obj = hashlib.new(algorithm)
+    hash_obj.update(string.encode())
+    return hash_obj.hexdigest()
+
+def hash_file_object(
+    file_obj: BinaryIO, algorithm: str = "sha256", chunk_size: int = 8192
+) -> str:
+    """
+    Calculate hash of a file-like object.
+
+    Args:
+        file_obj: File-like object (must be in binary mode)
+        algorithm: Hash algorithm (md5, sha1, sha256, etc.)
+        chunk_size: Size of chunks to read
+
+    Returns:
+        str: Hexadecimal digest of the hash
+    """
+    hash_obj = hashlib.new(algorithm)
+
+    # Store current position
+    try:
+        original_position = file_obj.tell()
+    except (OSError, IOError, AttributeError):
+        original_position = None
+
+    # Seek to the beginning if possible
+    try:
+        file_obj.seek(0)
+    except (OSError, IOError, AttributeError):
+        pass
+
+    # Calculate hash
+    for chunk in iter(lambda: file_obj.read(chunk_size), b""):
+        hash_obj.update(chunk)
+
+    # Restore original position if possible
+    if original_position is not None:
+        try:
+            file_obj.seek(original_position)
+        except (OSError, IOError, AttributeError):
+            pass
+
+    return hash_obj.hexdigest()
+