misata 0.2.0b0__py3-none-any.whl → 0.3.0b0__py3-none-any.whl

misata/context.py

@@ -0,0 +1,259 @@
+"""
+Context management for Misata data generation.
+
+Provides stateful context tracking during multi-table generation,
+including parent ID tracking for foreign keys and cross-table references.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Set
+
+import numpy as np
+import pandas as pd
+
+
+@dataclass
+class TableContext:
+    """Context for a single generated table."""
+    
+    name: str
+    row_count: int = 0
+    columns: Set[str] = field(default_factory=set)
+    primary_key: Optional[np.ndarray] = None
+    foreign_keys: Dict[str, np.ndarray] = field(default_factory=dict)
+    cached_columns: Dict[str, np.ndarray] = field(default_factory=dict)
+    
+    def set_primary_key(self, values: np.ndarray) -> None:
+        """Store primary key values for foreign key lookups."""
+        self.primary_key = values
+        self.row_count = len(values)
+    
+    def set_column(self, column_name: str, values: np.ndarray) -> None:
+        """Cache a column for cross-table references."""
+        self.cached_columns[column_name] = values
+        self.columns.add(column_name)
+    
+    def get_column(self, column_name: str) -> Optional[np.ndarray]:
+        """Get cached column values."""
+        return self.cached_columns.get(column_name)
+    
+    def get_ids(self) -> Optional[np.ndarray]:
+        """Get primary key values."""
+        return self.primary_key
+
+
+class GenerationContext:
+    """Manages state across multi-table data generation.
+    
+    This context tracks:
+    - Generated table data for foreign key references
+    - Columns needed for cross-table lookups
+    - Progress tracking for callbacks
+    
+    Example:
+        context = GenerationContext()
+        
+        # After generating users table
+        context.register_table("users", users_df)
+        
+        # When generating orders (which references users)
+        user_ids = context.get_parent_ids("users", "id")
+        orders_df["user_id"] = np.random.choice(user_ids, size=1000)
+    """
+    
+    def __init__(self):
+        self._tables: Dict[str, TableContext] = {}
+        self._generation_order: List[str] = []
+        self._progress_callbacks: List[callable] = []
+        self._current_table: Optional[str] = None
+        self._current_progress: float = 0.0
+    
+    def register_table(
+        self,
+        table_name: str,
+        df: pd.DataFrame,
+        id_column: str = "id"
+    ) -> None:
+        """Register a generated table in the context.
+        
+        Args:
+            table_name: Name of the table
+            df: Generated DataFrame
+            id_column: Primary key column name
+        """
+        ctx = TableContext(name=table_name)
+        
+        if id_column in df.columns:
+            ctx.set_primary_key(df[id_column].values)
+        
+        # Cache all columns for potential cross-references
+        for col in df.columns:
+            ctx.set_column(col, df[col].values)
+        
+        self._tables[table_name] = ctx
+        self._generation_order.append(table_name)
+    
+    def register_batch(
+        self,
+        table_name: str,
+        df: pd.DataFrame,
+        id_column: str = "id"
+    ) -> None:
+        """Register a batch of generated data (appends to existing).
+        
+        Args:
+            table_name: Name of the table
+            df: Generated batch DataFrame
+            id_column: Primary key column name
+        """
+        if table_name not in self._tables:
+            self.register_table(table_name, df, id_column)
+            return
+        
+        ctx = self._tables[table_name]
+        
+        # Append to existing
+        if id_column in df.columns:
+            if ctx.primary_key is not None:
+                ctx.primary_key = np.concatenate([ctx.primary_key, df[id_column].values])
+            else:
+                ctx.set_primary_key(df[id_column].values)
+        
+        for col in df.columns:
+            if col in ctx.cached_columns:
+                ctx.cached_columns[col] = np.concatenate([
+                    ctx.cached_columns[col],
+                    df[col].values
+                ])
+            else:
+                ctx.set_column(col, df[col].values)
+        
+        ctx.row_count = len(ctx.primary_key) if ctx.primary_key is not None else ctx.row_count + len(df)
+    
+    def get_parent_ids(
+        self,
+        table_name: str,
+        column: str = "id"
+    ) -> Optional[np.ndarray]:
+        """Get column values from a parent table for foreign key generation.
+        
+        Args:
+            table_name: Parent table name
+            column: Column to get values from
+            
+        Returns:
+            Array of values or None if table not found
+        """
+        if table_name not in self._tables:
+            return None
+        
+        ctx = self._tables[table_name]
+        
+        if column == "id" and ctx.primary_key is not None:
+            return ctx.primary_key
+        
+        return ctx.get_column(column)
+    
+    def get_filtered_parent_ids(
+        self,
+        table_name: str,
+        id_column: str = "id",
+        filters: Optional[Dict[str, Any]] = None
+    ) -> Optional[np.ndarray]:
+        """Get filtered parent IDs based on conditions.
+        
+        Args:
+            table_name: Parent table name
+            id_column: ID column to return
+            filters: Dict of column -> value conditions
+            
+        Returns:
+            Filtered array of IDs
+        """
+        if table_name not in self._tables:
+            return None
+        
+        ctx = self._tables[table_name]
+        
+        if not filters:
+            return self.get_parent_ids(table_name, id_column)
+        
+        # Get base IDs
+        ids = ctx.get_column(id_column)
+        if ids is None:
+            ids = ctx.primary_key
+        
+        if ids is None:
+            return None
+        
+        # Apply filters
+        mask = np.ones(len(ids), dtype=bool)
+        
+        for filter_col, filter_val in filters.items():
+            col_values = ctx.get_column(filter_col)
+            if col_values is not None:
+                mask &= (col_values == filter_val)
+        
+        return ids[mask] if mask.any() else None
+    
+    def get_table_context(self, table_name: str) -> Optional[TableContext]:
+        """Get full context for a table."""
+        return self._tables.get(table_name)
+    
+    def has_table(self, table_name: str) -> bool:
+        """Check if a table has been generated."""
+        return table_name in self._tables
+    
+    def get_generated_tables(self) -> List[str]:
+        """Get list of generated tables in order."""
+        return self._generation_order.copy()
+    
+    def clear(self) -> None:
+        """Clear all context data."""
+        self._tables.clear()
+        self._generation_order.clear()
+        self._current_table = None
+        self._current_progress = 0.0
+    
+    # ============ Progress Tracking ============
+    
+    def add_progress_callback(self, callback: callable) -> None:
+        """Add a progress callback function.
+        
+        Callback signature: callback(table_name: str, progress: float, message: str)
+        """
+        self._progress_callbacks.append(callback)
+    
+    def set_current_table(self, table_name: str) -> None:
+        """Set the currently generating table."""
+        self._current_table = table_name
+        self._notify_progress(0.0, f"Starting {table_name}")
+    
+    def update_progress(self, progress: float, message: str = "") -> None:
+        """Update generation progress (0.0 to 1.0)."""
+        self._current_progress = progress
+        self._notify_progress(progress, message)
+    
+    def _notify_progress(self, progress: float, message: str) -> None:
+        """Notify all progress callbacks."""
+        for callback in self._progress_callbacks:
+            try:
+                callback(self._current_table, progress, message)
+            except Exception:
+                pass  # Don't let callback errors break generation
+    
+    # ============ Statistics ============
+    
+    def get_summary(self) -> Dict[str, Any]:
+        """Get summary of all generated data."""
+        return {
+            "tables": {
+                name: {
+                    "row_count": ctx.row_count,
+                    "columns": list(ctx.columns),
+                }
+                for name, ctx in self._tables.items()
+            },
+            "generation_order": self._generation_order,
+            "total_rows": sum(ctx.row_count for ctx in self._tables.values()),
+        }

misata/exceptions.py

@@ -0,0 +1,277 @@
+"""
+Custom exceptions for Misata.
+
+Provides a hierarchy of exception classes for better error handling
+and more informative error messages.
+"""
+
+from typing import Any, Dict, List, Optional
+
+
+class MisataError(Exception):
+    """Base exception for all Misata errors."""
+    
+    def __init__(self, message: str, details: Optional[Dict[str, Any]] = None):
+        self.message = message
+        self.details = details or {}
+        super().__init__(self.message)
+    
+    def __str__(self) -> str:
+        if self.details:
+            return f"{self.message} | Details: {self.details}"
+        return self.message
+
+
+# ============ Schema Errors ============
+
+class SchemaError(MisataError):
+    """Base class for schema-related errors."""
+    pass
+
+
+class SchemaValidationError(SchemaError):
+    """Raised when schema validation fails."""
+    
+    def __init__(
+        self,
+        message: str,
+        field: Optional[str] = None,
+        value: Optional[Any] = None,
+        suggestion: Optional[str] = None,
+    ):
+        self.field = field
+        self.value = value
+        self.suggestion = suggestion
+        details = {}
+        if field:
+            details["field"] = field
+        if value is not None:
+            details["value"] = str(value)[:100]  # Truncate long values
+        if suggestion:
+            details["suggestion"] = suggestion
+        super().__init__(message, details)
+    
+    def __str__(self) -> str:
+        msg = self.message
+        if self.field:
+            msg = f"[{self.field}] {msg}"
+        if self.suggestion:
+            msg += f"\n  → Suggestion: {self.suggestion}"
+        return msg
+
+
+class SchemaParseError(SchemaError):
+    """Raised when schema parsing fails (YAML, JSON, etc.)."""
+    
+    def __init__(self, message: str, source: Optional[str] = None, line: Optional[int] = None):
+        self.source = source
+        self.line = line
+        details = {}
+        if source:
+            details["source"] = source
+        if line:
+            details["line"] = line
+        super().__init__(message, details)
+
+
+class RelationshipError(SchemaError):
+    """Raised when relationship definition is invalid."""
+    
+    def __init__(
+        self,
+        message: str,
+        parent_table: Optional[str] = None,
+        child_table: Optional[str] = None,
+    ):
+        self.parent_table = parent_table
+        self.child_table = child_table
+        details = {}
+        if parent_table:
+            details["parent_table"] = parent_table
+        if child_table:
+            details["child_table"] = child_table
+        super().__init__(message, details)
+
+
+# ============ Generation Errors ============
+
+class GenerationError(MisataError):
+    """Base class for data generation errors."""
+    pass
+
+
+class ColumnGenerationError(GenerationError):
+    """Raised when column generation fails."""
+    
+    def __init__(
+        self,
+        message: str,
+        table: Optional[str] = None,
+        column: Optional[str] = None,
+        column_type: Optional[str] = None,
+        suggestion: Optional[str] = None,
+    ):
+        self.table = table
+        self.column = column
+        self.column_type = column_type
+        self.suggestion = suggestion
+        details = {}
+        if table:
+            details["table"] = table
+        if column:
+            details["column"] = column
+        if column_type:
+            details["column_type"] = column_type
+        super().__init__(message, details)
+    
+    def __str__(self) -> str:
+        location = ""
+        if self.table and self.column:
+            location = f"[{self.table}.{self.column}] "
+        elif self.table:
+            location = f"[{self.table}] "
+        msg = f"{location}{self.message}"
+        if self.suggestion:
+            msg += f"\n  → Suggestion: {self.suggestion}"
+        return msg
+
+
+class ConstraintError(GenerationError):
+    """Raised when constraint application fails."""
+    
+    def __init__(
+        self,
+        message: str,
+        constraint_type: Optional[str] = None,
+        affected_columns: Optional[List[str]] = None,
+    ):
+        self.constraint_type = constraint_type
+        self.affected_columns = affected_columns or []
+        details = {}
+        if constraint_type:
+            details["constraint_type"] = constraint_type
+        if affected_columns:
+            details["affected_columns"] = affected_columns
+        super().__init__(message, details)
+
+
+class CircularDependencyError(GenerationError):
+    """Raised when circular dependencies are detected in table relationships."""
+    
+    def __init__(self, message: str, tables: Optional[List[str]] = None):
+        self.tables = tables or []
+        details = {"tables": tables} if tables else {}
+        super().__init__(message, details)
+
+
+class ReferentialIntegrityError(GenerationError):
+    """Raised when foreign key references cannot be satisfied."""
+    
+    def __init__(
+        self,
+        message: str,
+        parent_table: Optional[str] = None,
+        child_table: Optional[str] = None,
+        missing_ids: Optional[int] = None,
+    ):
+        self.parent_table = parent_table
+        self.child_table = child_table
+        self.missing_ids = missing_ids
+        details = {}
+        if parent_table:
+            details["parent_table"] = parent_table
+        if child_table:
+            details["child_table"] = child_table
+        if missing_ids:
+            details["missing_ids"] = missing_ids
+        super().__init__(message, details)
+
+
+# ============ LLM Errors ============
+
+class LLMError(MisataError):
+    """Base class for LLM-related errors."""
+    pass
+
+
+class LLMConnectionError(LLMError):
+    """Raised when LLM API connection fails."""
+    
+    def __init__(self, message: str, provider: Optional[str] = None):
+        self.provider = provider
+        details = {"provider": provider} if provider else {}
+        super().__init__(message, details)
+
+
+class LLMParseError(LLMError):
+    """Raised when LLM response cannot be parsed."""
+    
+    def __init__(self, message: str, raw_response: Optional[str] = None):
+        self.raw_response = raw_response[:500] if raw_response else None
+        details = {}
+        if self.raw_response:
+            details["raw_response_preview"] = self.raw_response
+        super().__init__(message, details)
+
+
+class LLMQuotaError(LLMError):
+    """Raised when LLM API quota is exceeded."""
+    
+    def __init__(self, message: str, provider: Optional[str] = None, retry_after: Optional[int] = None):
+        self.provider = provider
+        self.retry_after = retry_after
+        details = {}
+        if provider:
+            details["provider"] = provider
+        if retry_after:
+            details["retry_after_seconds"] = retry_after
+        super().__init__(message, details)
+
+
+# ============ Configuration Errors ============
+
+class ConfigurationError(MisataError):
+    """Raised when configuration is invalid."""
+    pass
+
+
+class MissingAPIKeyError(ConfigurationError):
+    """Raised when required API key is missing."""
+    
+    def __init__(self, provider: str, env_var: str):
+        self.provider = provider
+        self.env_var = env_var
+        message = f"API key for {provider} not found"
+        suggestion = f"Set {env_var} environment variable or pass api_key parameter"
+        details = {"provider": provider, "env_var": env_var, "suggestion": suggestion}
+        super().__init__(message, details)
+
+
+# ============ Export Errors ============
+
+class ExportError(MisataError):
+    """Base class for export-related errors."""
+    pass
+
+
+class FileWriteError(ExportError):
+    """Raised when file writing fails."""
+    
+    def __init__(self, message: str, path: Optional[str] = None):
+        self.path = path
+        details = {"path": path} if path else {}
+        super().__init__(message, details)
+
+
+class InvalidOutputFormatError(ExportError):
+    """Raised when export format is not supported."""
+    
+    def __init__(self, format: str, supported_formats: Optional[List[str]] = None):
+        self.format = format
+        self.supported_formats = supported_formats or ["csv", "parquet", "json"]
+        message = f"Unsupported export format: {format}"
+        details = {
+            "format": format,
+            "supported_formats": self.supported_formats,
+        }
+        super().__init__(message, details)

misata/generators/__init__.py

@@ -0,0 +1,29 @@
+"""
+Generators package for Misata.
+
+Provides type-safe data generators for all supported column types.
+"""
+
+from misata.generators.base import (
+    BaseGenerator,
+    BooleanGenerator,
+    CategoricalGenerator,
+    DateGenerator,
+    FloatGenerator,
+    ForeignKeyGenerator,
+    GeneratorFactory,
+    IntegerGenerator,
+    TextGenerator,
+)
+
+__all__ = [
+    "BaseGenerator",
+    "GeneratorFactory",
+    "IntegerGenerator",
+    "FloatGenerator",
+    "BooleanGenerator",
+    "CategoricalGenerator",
+    "DateGenerator",
+    "TextGenerator",
+    "ForeignKeyGenerator",
+]