additory 0.1.0a3__py3-none-any.whl → 0.1.1a1__py3-none-any.whl

additory/core/memory_manager.py

@@ -1,547 +1,209 @@
-# memory_manager.py
-# Centralized memory management for enhanced expressions system
+"""
+Memory management for Additory operations.
 
-import gc
-import os
-import psutil
-import threading
-import time
-from typing import Dict, List, Optional, Callable, Any
-from dataclasses import dataclass, field
-from datetime import datetime, timedelta
-from contextlib import contextmanager
-
-from .logging import log_info, log_warning
-
-
-@dataclass
-class MemorySnapshot:
-    """Memory usage snapshot at a point in time"""
-    timestamp: datetime
-    process_memory_mb: float
-    arrow_allocated_bytes: int
-    python_objects_count: int
-    gc_collections: Dict[int, int]
-    custom_metrics: Dict[str, Any] = field(default_factory=dict)
+Provides automatic cleanup of intermediate objects to free memory:
+- Tracks temporary DataFrames
+- Cleans up after operations
+- Can be enabled/disabled for debugging
+"""
 
-
-@dataclass
-class MemoryThresholds:
-    """Memory usage thresholds for cleanup triggers"""
-    process_memory_mb: float = 500.0  # Process memory threshold
-    arrow_memory_bytes: int = 100 * 1024 * 1024  # 100MB Arrow memory
-    python_objects_count: int = 100000  # Python object count
-    cleanup_interval_seconds: float = 30.0  # Periodic cleanup interval
-
-
-class MemoryLeakDetector:
-    """Detects potential memory leaks by tracking memory growth patterns"""
-    
-    def __init__(self, window_size: int = 10, growth_threshold: float = 1.5):
-        self.window_size = window_size
-        self.growth_threshold = growth_threshold
-        self.snapshots: List[MemorySnapshot] = []
-        self.leak_warnings = 0
-    
-    def add_snapshot(self, snapshot: MemorySnapshot):
-        """Add a memory snapshot for leak detection"""
-        self.snapshots.append(snapshot)
-        
-        # Keep only the last window_size snapshots
-        if len(self.snapshots) > self.window_size:
-            self.snapshots.pop(0)
-        
-        # Check for potential leaks
-        if len(self.snapshots) >= self.window_size:
-            self._check_for_leaks()
-    
-    def _check_for_leaks(self):
-        """Check if memory usage shows signs of leaking"""
-        if len(self.snapshots) < 2:
-            return
-        
-        # Calculate memory growth rate
-        first_snapshot = self.snapshots[0]
-        last_snapshot = self.snapshots[-1]
-        
-        memory_growth = last_snapshot.process_memory_mb / max(first_snapshot.process_memory_mb, 1.0)  # Avoid division by zero
-        time_span = (last_snapshot.timestamp - first_snapshot.timestamp).total_seconds()
-        
-        # Check for sustained memory growth (relaxed time threshold for testing)
-        if memory_growth > self.growth_threshold and time_span > 10:  # 10 seconds minimum
-            self.leak_warnings += 1
-            log_warning(f"[memory_manager] Potential memory leak detected: "
-                       f"{memory_growth:.2f}x growth over {time_span:.1f}s")
-    
-    def get_leak_status(self) -> Dict[str, Any]:
-        """Get current leak detection status"""
-        return {
-            "leak_warnings": self.leak_warnings,
-            "snapshots_count": len(self.snapshots),
-            "monitoring_window": self.window_size,
-            "growth_threshold": self.growth_threshold
-        }
+import gc
+import sys
+from typing import Any, Callable, List, Optional, Tuple
+from functools import wraps
 
 
 class MemoryManager:
-    """Centralized memory management for enhanced expressions system"""
+    """
+    Manages memory cleanup after operations.
+    
+    Tracks intermediate objects created during processing and
+    cleans them up automatically to free memory.
+    """
     
     def __init__(self):
-        # Configuration
-        self.thresholds = MemoryThresholds()
-        self.monitoring_enabled = True
-        self.auto_cleanup_enabled = True
-        
-        # State tracking
-        self.snapshots: List[MemorySnapshot] = []
-        self.cleanup_callbacks: List[Callable[[], None]] = []
-        self.leak_detector = MemoryLeakDetector()
-        
-        # Statistics
-        self.stats = {
-            "total_cleanups": 0,
-            "forced_cleanups": 0,
-            "auto_cleanups": 0,
-            "memory_freed_mb": 0.0,
-            "last_cleanup": None
-        }
-        
-        # Background monitoring
-        self._monitoring_thread = None
-        self._stop_monitoring = threading.Event()
-        
-        # Process reference for memory monitoring
-        try:
-            self.process = psutil.Process(os.getpid())
-        except Exception as e:
-            log_warning(f"[memory_manager] Failed to initialize process monitor: {e}")
-            self.process = None
-        
-        log_info("[memory_manager] Memory Manager initialized")
+        """Initialize memory manager."""
+        self.tracked_objects: List[Tuple[Any, Optional[str]]] = []
+        self.cleanup_enabled: bool = True
     
-    def start_monitoring(self, interval_seconds: float = 30.0):
-        """Start background memory monitoring"""
-        if self._monitoring_thread and self._monitoring_thread.is_alive():
-            log_warning("[memory_manager] Monitoring already running")
-            return
-        
-        self.thresholds.cleanup_interval_seconds = interval_seconds
-        self._stop_monitoring.clear()
-        
-        self._monitoring_thread = threading.Thread(
-            target=self._monitoring_loop,
-            daemon=True,
-            name="MemoryMonitor"
-        )
-        self._monitoring_thread.start()
+    def track(self, obj: Any, name: Optional[str] = None) -> None:
+        """
+        Track object for cleanup.
         
-        log_info(f"[memory_manager] Started background monitoring (interval: {interval_seconds}s)")
-    
-    def stop_monitoring(self):
-        """Stop background memory monitoring"""
-        if self._monitoring_thread and self._monitoring_thread.is_alive():
-            self._stop_monitoring.set()
-            self._monitoring_thread.join(timeout=5.0)
-            log_info("[memory_manager] Stopped background monitoring")
-    
-    def _monitoring_loop(self):
-        """Background monitoring loop"""
-        while not self._stop_monitoring.is_set():
-            try:
-                if self.monitoring_enabled:
-                    # Take memory snapshot
-                    snapshot = self.take_snapshot()
-                    
-                    # Check for cleanup needs
-                    if self.auto_cleanup_enabled and self._should_cleanup(snapshot):
-                        self.cleanup_if_needed()
-                
-                # Wait for next interval
-                self._stop_monitoring.wait(self.thresholds.cleanup_interval_seconds)
-                
-            except Exception as e:
-                log_warning(f"[memory_manager] Monitoring loop error: {e}")
-                time.sleep(5)  # Brief pause before retrying
-    
-    def register_cleanup_callback(self, callback: Callable[[], None]):
-        """Register a callback function for memory cleanup"""
-        self.cleanup_callbacks.append(callback)
-        log_info(f"[memory_manager] Registered cleanup callback: {callback.__name__}")
-    
-    def unregister_cleanup_callback(self, callback: Callable[[], None]):
-        """Unregister a cleanup callback"""
-        if callback in self.cleanup_callbacks:
-            self.cleanup_callbacks.remove(callback)
-            log_info(f"[memory_manager] Unregistered cleanup callback: {callback.__name__}")
-    
-    def take_snapshot(self) -> MemorySnapshot:
-        """Take a snapshot of current memory usage"""
-        try:
-            # Process memory
-            process_memory_mb = 0.0
-            if self.process:
-                memory_info = self.process.memory_info()
-                process_memory_mb = memory_info.rss / 1024 / 1024
-            
-            # Arrow memory (if available)
-            arrow_allocated_bytes = 0
-            try:
-                import pyarrow as pa
-                arrow_allocated_bytes = pa.default_memory_pool().bytes_allocated()
-            except ImportError:
-                pass
-            
-            # Python objects
-            python_objects_count = len(gc.get_objects())
-            
-            # GC statistics
-            gc_collections = {i: gc.get_count()[i] for i in range(3)}
-            
-            snapshot = MemorySnapshot(
-                timestamp=datetime.now(),
-                process_memory_mb=process_memory_mb,
-                arrow_allocated_bytes=arrow_allocated_bytes,
-                python_objects_count=python_objects_count,
-                gc_collections=gc_collections
-            )
-            
-            # Store snapshot
-            self.snapshots.append(snapshot)
-            
-            # Keep only recent snapshots (last 100)
-            if len(self.snapshots) > 100:
-                self.snapshots.pop(0)
-            
-            # Add to leak detector
-            self.leak_detector.add_snapshot(snapshot)
+        Args:
+            obj: Object to track (typically DataFrame or Series)
+            name: Optional name for debugging
             
-            return snapshot
-            
-        except Exception as e:
-            log_warning(f"[memory_manager] Failed to take memory snapshot: {e}")
-            return MemorySnapshot(
-                timestamp=datetime.now(),
-                process_memory_mb=0.0,
-                arrow_allocated_bytes=0,
-                python_objects_count=0,
-                gc_collections={}
-            )
-    
-    def _should_cleanup(self, snapshot: MemorySnapshot) -> bool:
-        """Check if cleanup is needed based on thresholds"""
-        return (
-            snapshot.process_memory_mb > self.thresholds.process_memory_mb or
-            snapshot.arrow_allocated_bytes > self.thresholds.arrow_memory_bytes or
-            snapshot.python_objects_count > self.thresholds.python_objects_count
-        )
+        Example:
+            memory_manager.track(temp_df, 'intermediate_result')
+        """
+        if self.cleanup_enabled:
+            self.tracked_objects.append((obj, name))
     
-    def cleanup_if_needed(self) -> bool:
-        """Perform cleanup if memory usage exceeds thresholds"""
-        snapshot = self.take_snapshot()
-        
-        if not self._should_cleanup(snapshot):
-            return False
-        
-        log_info(f"[memory_manager] Auto cleanup triggered - "
-                f"Memory: {snapshot.process_memory_mb:.1f}MB, "
-                f"Arrow: {snapshot.arrow_allocated_bytes} bytes, "
-                f"Objects: {snapshot.python_objects_count}")
-        
-        memory_before = snapshot.process_memory_mb
-        self._perform_cleanup()
-        
-        # Take another snapshot to measure cleanup effectiveness
-        after_snapshot = self.take_snapshot()
-        memory_freed = max(0, memory_before - after_snapshot.process_memory_mb)
-        
-        self.stats["auto_cleanups"] += 1
-        self.stats["memory_freed_mb"] += memory_freed
-        self.stats["last_cleanup"] = datetime.now()
-        
-        log_info(f"[memory_manager] Auto cleanup completed - "
-                f"Freed: {memory_freed:.1f}MB")
-        
-        return True
-    
-    def force_cleanup(self) -> Dict[str, float]:
-        """Force immediate memory cleanup"""
-        log_info("[memory_manager] Forcing memory cleanup")
-        
-        before_snapshot = self.take_snapshot()
-        memory_before = before_snapshot.process_memory_mb
-        
-        self._perform_cleanup()
-        
-        after_snapshot = self.take_snapshot()
-        memory_after = after_snapshot.process_memory_mb
-        memory_freed = max(0, memory_before - memory_after)
-        
-        self.stats["forced_cleanups"] += 1
-        self.stats["memory_freed_mb"] += memory_freed
-        self.stats["last_cleanup"] = datetime.now()
-        
-        cleanup_stats = {
-            "memory_before_mb": memory_before,
-            "memory_after_mb": memory_after,
-            "memory_freed_mb": memory_freed,
-            "arrow_freed_bytes": before_snapshot.arrow_allocated_bytes - after_snapshot.arrow_allocated_bytes
-        }
+    def cleanup(self) -> None:
+        """
+        Clean up all tracked objects.
         
-        log_info(f"[memory_manager] Forced cleanup completed - "
-                f"Freed: {memory_freed:.1f}MB")
+        Deletes references to tracked objects and runs garbage collection.
+        Called automatically at the end of every operation.
         
-        return cleanup_stats
-    
-    def _perform_cleanup(self):
-        """Perform the actual cleanup operations"""
-        # Call registered cleanup callbacks
-        for callback in self.cleanup_callbacks:
+        Example:
             try:
-                callback()
-            except Exception as e:
-                log_warning(f"[memory_manager] Cleanup callback failed: {e}")
-        
-        # Force garbage collection
-        collected = gc.collect()
-        
-        # Additional cleanup for specific libraries
-        self._cleanup_arrow_memory()
-        self._cleanup_polars_memory()
+                result = process_data(df)
+                return result
+            finally:
+                memory_manager.cleanup()
+        """
+        if not self.cleanup_enabled:
+            return
         
-        self.stats["total_cleanups"] += 1
+        # Count objects before cleanup
+        count = len(self.tracked_objects)
         
-        log_info(f"[memory_manager] Cleanup performed - "
-                f"GC collected: {collected} objects")
-    
-    def _cleanup_arrow_memory(self):
-        """Cleanup Arrow-specific memory"""
-        try:
-            import pyarrow as pa
-            # Force Arrow memory pool cleanup
-            pool = pa.default_memory_pool()
-            allocated_before = pool.bytes_allocated()
-            
-            # Trigger garbage collection to free Arrow objects
-            gc.collect()
+        if count > 0:
+            # Estimate memory before cleanup (optional, for logging)
+            total_memory = 0
+            for obj, name in self.tracked_objects:
+                try:
+                    memory = estimate_memory_usage(obj)
+                    total_memory += memory
+                except:
+                    pass  # Ignore errors in memory estimation
             
-            allocated_after = pool.bytes_allocated()
-            freed = allocated_before - allocated_after
+            # Clear references
+            self.tracked_objects.clear()
             
-            if freed > 0:
-                log_info(f"[memory_manager] Arrow cleanup freed {freed} bytes")
-                
-        except ImportError:
-            pass
-        except Exception as e:
-            log_warning(f"[memory_manager] Arrow cleanup failed: {e}")
-    
-    def _cleanup_polars_memory(self):
-        """Cleanup Polars-specific memory"""
-        try:
-            import polars as pl
-            # Polars cleanup is mostly handled by Rust's memory management
-            # But we can clear any cached data
+            # Force garbage collection
             gc.collect()
             
-        except ImportError:
-            pass
-        except Exception as e:
-            log_warning(f"[memory_manager] Polars cleanup failed: {e}")
-    
-    def get_memory_stats(self) -> Dict[str, Any]:
-        """Get comprehensive memory statistics"""
-        current_snapshot = self.take_snapshot()
-        
-        # Calculate memory trends
-        memory_trend = "stable"
-        if len(self.snapshots) >= 10:
-            recent_memory = sum(s.process_memory_mb for s in self.snapshots[-5:]) / min(5, len(self.snapshots))
-            older_memory = sum(s.process_memory_mb for s in self.snapshots[-10:-5]) / max(1, min(5, len(self.snapshots) - 5))
-            
-            if recent_memory > older_memory * 1.1:
-                memory_trend = "increasing"
-            elif recent_memory < older_memory * 0.9:
-                memory_trend = "decreasing"
-        elif len(self.snapshots) >= 2:
-            # Simple trend for fewer snapshots
-            if self.snapshots[-1].process_memory_mb > self.snapshots[0].process_memory_mb * 1.1:
-                memory_trend = "increasing"
-            elif self.snapshots[-1].process_memory_mb < self.snapshots[0].process_memory_mb * 0.9:
-                memory_trend = "decreasing"
-        
-        return {
-            "current": {
-                "process_memory_mb": current_snapshot.process_memory_mb,
-                "arrow_allocated_bytes": current_snapshot.arrow_allocated_bytes,
-                "python_objects_count": current_snapshot.python_objects_count,
-                "gc_collections": current_snapshot.gc_collections
-            },
-            "thresholds": {
-                "process_memory_mb": self.thresholds.process_memory_mb,
-                "arrow_memory_bytes": self.thresholds.arrow_memory_bytes,
-                "python_objects_count": self.thresholds.python_objects_count
-            },
-            "trends": {
-                "memory_trend": memory_trend,
-                "snapshots_count": len(self.snapshots),
-                "monitoring_enabled": self.monitoring_enabled
-            },
-            "cleanup_stats": self.stats.copy(),
-            "leak_detection": self.leak_detector.get_leak_status()
-        }
-    
-    def get_memory_usage_mb(self) -> float:
-        """Get current memory usage in MB"""
-        if self.process:
+            # Log cleanup (if logger is available)
             try:
-                return self.process.memory_info().rss / 1024 / 1024
-            except Exception:
-                pass
-        return 0.0
-    
-    def set_thresholds(self, **kwargs):
-        """Update memory thresholds"""
-        for key, value in kwargs.items():
-            if hasattr(self.thresholds, key):
-                setattr(self.thresholds, key, value)
-                log_info(f"[memory_manager] Updated threshold {key} = {value}")
-            else:
-                log_warning(f"[memory_manager] Unknown threshold: {key}")
-    
-    def enable_monitoring(self):
-        """Enable memory monitoring"""
-        self.monitoring_enabled = True
-        log_info("[memory_manager] Memory monitoring enabled")
+                from .logging import get_logger
+                logger = get_logger()
+                if total_memory > 0:
+                    memory_mb = total_memory / (1024 * 1024)
+                    logger.info(
+                        f"Cleaned up {count} intermediate objects (~{memory_mb:.2f} MB)"
+                    )
+            except:
+                pass  # Ignore if logger not available
     
-    def disable_monitoring(self):
-        """Disable memory monitoring"""
-        self.monitoring_enabled = False
-        log_info("[memory_manager] Memory monitoring disabled")
-    
-    def enable_auto_cleanup(self):
-        """Enable automatic cleanup"""
-        self.auto_cleanup_enabled = True
-        log_info("[memory_manager] Auto cleanup enabled")
-    
-    def disable_auto_cleanup(self):
-        """Disable automatic cleanup"""
-        self.auto_cleanup_enabled = False
-        log_info("[memory_manager] Auto cleanup disabled")
+    def enable(self) -> None:
+        """
+        Enable automatic cleanup.
+        
+        Called by user code via add.enable_memory_cleanup()
+        """
+        self.cleanup_enabled = True
     
-    @contextmanager
-    def memory_context(self, description: str = "operation"):
-        """Context manager for monitoring memory usage during operations"""
-        before_snapshot = self.take_snapshot()
-        start_time = datetime.now()
+    def disable(self) -> None:
+        """
+        Disable automatic cleanup.
         
-        log_info(f"[memory_manager] Starting {description} - "
-                f"Memory: {before_snapshot.process_memory_mb:.1f}MB")
+        Useful for debugging to inspect intermediate objects.
+        Called by user code via add.disable_memory_cleanup()
+        """
+        self.cleanup_enabled = False
+    
+    def get_tracked_count(self) -> int:
+        """
+        Get number of tracked objects.
         
-        try:
-            yield self
-        finally:
-            after_snapshot = self.take_snapshot()
-            duration = (datetime.now() - start_time).total_seconds()
-            memory_delta = after_snapshot.process_memory_mb - before_snapshot.process_memory_mb
+        Returns:
+            Number of tracked objects
             
-            log_info(f"[memory_manager] Completed {description} - "
-                    f"Duration: {duration:.2f}s, "
-                    f"Memory delta: {memory_delta:+.1f}MB")
-            
-            # Auto cleanup if needed
-            if self.auto_cleanup_enabled and memory_delta > 50:  # 50MB increase
-                self.cleanup_if_needed()
-    
-    def reset_stats(self):
-        """Reset memory management statistics"""
-        self.stats = {
-            "total_cleanups": 0,
-            "forced_cleanups": 0,
-            "auto_cleanups": 0,
-            "memory_freed_mb": 0.0,
-            "last_cleanup": None
-        }
-        self.leak_detector = MemoryLeakDetector()
-        log_info("[memory_manager] Statistics reset")
+        Used for testing and debugging.
+        """
+        return len(self.tracked_objects)
     
-    def __del__(self):
-        """Cleanup when memory manager is destroyed"""
-        try:
-            self.stop_monitoring()
-        except Exception:
-            pass
+    def clear_tracking(self) -> None:
+        """
+        Clear tracking list without cleanup.
+        
+        Called after successful cleanup or when resetting state.
+        """
+        self.tracked_objects.clear()
 
 
 # Global memory manager instance
-_global_memory_manager = None
+_global_memory_manager: Optional[MemoryManager] = None
 
 
 def get_memory_manager() -> MemoryManager:
-    """Get the global memory manager instance"""
+    """
+    Get the global memory manager instance.
+    
+    Returns:
+        Global MemoryManager instance
+        
+    Example:
+        memory_manager = get_memory_manager()
+        memory_manager.track(temp_df)
+    """
     global _global_memory_manager
     if _global_memory_manager is None:
         _global_memory_manager = MemoryManager()
     return _global_memory_manager
 
 
-def cleanup_memory():
-    """Convenience function for forcing memory cleanup"""
-    return get_memory_manager().force_cleanup()
-
-
-def get_memory_stats():
-    """Convenience function for getting memory statistics"""
-    return get_memory_manager().get_memory_stats()
-
-
-def memory_context(description: str = "operation"):
-    """Convenience function for memory monitoring context"""
-    return get_memory_manager().memory_context(description)
-
-
-def memory_profile(description: str = None):
+def cleanup_after_operation(func: Callable) -> Callable:
     """
-    Decorator for profiling memory usage of functions.
+    Decorator to ensure cleanup after operation.
+    
+    Wraps a function to automatically clean up tracked objects
+    after execution, even if an exception occurs.
     
     Args:
-        description: Optional description for the operation
+        func: Function to wrap
+        
+    Returns:
+        Wrapped function with automatic cleanup
         
     Example:
-        @memory_profile("data processing")
-        def process_data(df):
-            return df.with_columns(...)
-    """
-    def decorator(func):
-        def wrapper(*args, **kwargs):
-            func_description = description or f"{func.__module__}.{func.__name__}"
-            with memory_context(func_description):
-                return func(*args, **kwargs)
-        return wrapper
-    return decorator
-
-
-def track_memory_usage(func):
-    """
-    Simple decorator to track memory usage of a function.
-    
-    Example:
-        @track_memory_usage
-        def expensive_operation():
-            # ... do work
+        @cleanup_after_operation
+        def to(df, ...):
+            # Function implementation
             pass
     """
+    @wraps(func)
     def wrapper(*args, **kwargs):
-        manager = get_memory_manager()
-        before = manager.get_memory_usage_mb()
-        
+        memory_manager = get_memory_manager()
         try:
             result = func(*args, **kwargs)
             return result
         finally:
-            after = manager.get_memory_usage_mb()
-            delta = after - before
-            
-            from .logging import log_info
-            log_info(f"[memory_profile] {func.__name__}: {delta:+.1f}MB")
+            memory_manager.cleanup()
+    
+    return wrapper
+
+
+def estimate_memory_usage(obj: Any) -> int:
+    """
+    Estimate memory usage of an object.
     
-    return wrapper
+    Args:
+        obj: Object to estimate (typically DataFrame or Series)
+        
+    Returns:
+        Estimated memory usage in bytes
+        
+    Note:
+        This is an approximation. Actual memory usage may vary.
+    """
+    try:
+        # Try to get memory usage from object (Polars/pandas)
+        if hasattr(obj, 'estimated_size'):
+            # Polars DataFrame
+            return obj.estimated_size()
+        elif hasattr(obj, 'memory_usage'):
+            # pandas DataFrame or Series
+            memory_usage = obj.memory_usage(deep=True)
+            if hasattr(memory_usage, 'sum'):
+                return int(memory_usage.sum())
+            return int(memory_usage)
+        else:
+            # Fallback to sys.getsizeof
+            return sys.getsizeof(obj)
+    except:
+        # If all else fails, return 0
+        return 0