oscura 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

oscura/utils/memory_extensions.py

@@ -0,0 +1,465 @@
+"""Extended memory management utilities for TraceKit.
+
+Additional memory management features including context managers,
+LRU caching, and HDF5 lazy loading support.
+
+
+Example:
+    >>> from oscura.utils.memory_extensions import ResourceManager, LRUCache
+    >>> with ResourceManager(large_array) as data:
+    ...     result = process(data)
+    >>> # Data automatically cleaned up
+
+References:
+    Python resource management patterns
+    functools.lru_cache documentation
+"""
+
+from __future__ import annotations
+
+import gc
+import hashlib
+import os
+import time
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, TypeVar
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from numpy.typing import NDArray
+
+T = TypeVar("T")
+
+
+# =============================================================================
+# Context Managers for Resource Cleanup (MEM-019)
+# =============================================================================
+
+
+class ResourceManager:
+    """Context manager for large data resources with automatic cleanup.
+
+    Ensures prompt memory release when done with large datasets.
+
+    Args:
+        resource: Resource to manage (array, file handle, etc.).
+        cleanup_func: Optional cleanup function to call on exit.
+
+    Example:
+        >>> import numpy as np
+        >>> with ResourceManager(np.zeros(1000000)) as data:
+        ...     result = process(data)
+        >>> # Data is automatically released
+
+    References:
+        MEM-019: Explicit Resource Cleanup
+    """
+
+    def __init__(
+        self,
+        resource: Any,
+        cleanup_func: Callable[[Any], None] | None = None,
+    ) -> None:
+        self._resource = resource
+        self._cleanup_func = cleanup_func
+
+    def __enter__(self) -> Any:
+        """Enter context, return resource."""
+        return self._resource
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit context, cleanup resource."""
+        # Note: exc_val and exc_tb intentionally unused but required for Python 3.11+ compatibility
+        if self._cleanup_func is not None:
+            self._cleanup_func(self._resource)
+
+        # Delete reference
+        self._resource = None
+
+        # Force garbage collection
+        gc.collect()
+
+
+class ArrayManager(ResourceManager):
+    """Context manager specifically for numpy arrays.
+
+    Example:
+        >>> import numpy as np
+        >>> with ArrayManager(np.zeros((10000, 10000))) as arr:
+        ...     result = np.sum(arr)
+        >>> # Array memory is released
+    """
+
+    def __init__(self, array: NDArray[Any]) -> None:
+        super().__init__(array, cleanup_func=lambda x: None)  # noqa: ARG005
+
+
+# =============================================================================
+# LRU Cache for Intermediate Results (MEM-021, MEM-029)
+# =============================================================================
+
+
+class LRUCache[T]:
+    """Least-Recently-Used cache with memory-based eviction.
+
+    Caches intermediate results with automatic eviction when
+    memory limit is exceeded.
+
+    Args:
+        max_memory_bytes: Maximum cache size in bytes.
+        max_entries: Maximum number of cache entries (default unlimited).
+
+    Example:
+        >>> cache = LRUCache(max_memory_bytes=1_000_000_000)  # 1 GB
+        >>> cache.put("key1", large_array, size_bytes=800_000_000)
+        >>> result = cache.get("key1")
+        >>> cache.clear()
+
+    References:
+        MEM-021: Intermediate Result Eviction
+        MEM-029: LRU Cache for Intermediate Results
+    """
+
+    def __init__(
+        self,
+        max_memory_bytes: int,
+        max_entries: int | None = None,
+    ) -> None:
+        self._max_memory = max_memory_bytes
+        self._max_entries = max_entries
+        self._cache: OrderedDict[str, tuple[T, int, float]] = OrderedDict()
+        self._current_memory: int = 0
+        self._hits: int = 0
+        self._misses: int = 0
+
+    def get(self, key: str) -> T | None:
+        """Get cached value by key.
+
+        Args:
+            key: Cache key.
+
+        Returns:
+            Cached value or None if not found.
+        """
+        if key in self._cache:
+            # Move to end (most recently used)
+            value, size, _ = self._cache.pop(key)
+            self._cache[key] = (value, size, time.time())
+            self._hits += 1
+            return value
+        else:
+            self._misses += 1
+            return None
+
+    def put(self, key: str, value: T, size_bytes: int | None = None) -> None:
+        """Put value in cache.
+
+        Args:
+            key: Cache key.
+            value: Value to cache.
+            size_bytes: Size in bytes. If None, estimated from value.
+        """
+        # Estimate size if not provided
+        if size_bytes is None:
+            size_bytes = self._estimate_size(value)
+
+        # Check if single item exceeds max memory
+        if size_bytes > self._max_memory:
+            # Don't cache items larger than max memory
+            return
+
+        # Evict until we have space
+        while (
+            self._current_memory + size_bytes > self._max_memory
+            or (self._max_entries and len(self._cache) >= self._max_entries)
+        ) and len(self._cache) > 0:
+            self._evict_oldest()
+
+        # Remove if key already exists
+        if key in self._cache:
+            _, old_size, _ = self._cache.pop(key)
+            self._current_memory -= old_size
+
+        # Add new entry
+        self._cache[key] = (value, size_bytes, time.time())
+        self._current_memory += size_bytes
+
+    def _evict_oldest(self) -> None:
+        """Evict least recently used entry."""
+        if len(self._cache) > 0:
+            _, (_, size, _) = self._cache.popitem(last=False)
+            self._current_memory -= size
+
+    def clear(self) -> None:
+        """Clear entire cache."""
+        self._cache.clear()
+        self._current_memory = 0
+
+    def _estimate_size(self, value: T) -> int:
+        """Estimate size of cached value."""
+        if isinstance(value, np.ndarray):
+            return int(value.nbytes)
+        elif isinstance(value, list | tuple):
+            # Rough estimate for sequences
+            return sum(self._estimate_size(item) for item in value)
+        elif isinstance(value, dict):
+            return sum(self._estimate_size(k) + self._estimate_size(v) for k, v in value.items())
+        else:
+            # Fallback: assume 1KB for unknown types
+            return 1024
+
+    def stats(self) -> dict[str, int | float]:
+        """Get cache statistics.
+
+        Returns:
+            Dictionary with cache stats.
+        """
+        total_requests = self._hits + self._misses
+        hit_rate = self._hits / total_requests if total_requests > 0 else 0.0
+
+        return {
+            "entries": len(self._cache),
+            "memory_bytes": self._current_memory,
+            "memory_mb": self._current_memory / (1024 * 1024),
+            "hits": self._hits,
+            "misses": self._misses,
+            "hit_rate": hit_rate,
+        }
+
+    def __len__(self) -> int:
+        """Number of cached entries."""
+        return len(self._cache)
+
+
+# Global result cache
+_result_cache: LRUCache[Any] | None = None
+
+
+def get_result_cache() -> LRUCache[Any]:
+    """Get global result cache.
+
+    Returns:
+        Global LRU cache instance.
+
+    Example:
+        >>> cache = get_result_cache()
+        >>> cache.put("fft_result", fft_data, size_bytes=8000000)
+        >>> result = cache.get("fft_result")
+    """
+    global _result_cache
+    if _result_cache is None:
+        # Default: 2 GB cache
+        max_cache_size = int(os.environ.get("TK_CACHE_SIZE", 2 * 1024 * 1024 * 1024))  # noqa: PLW1508
+        _result_cache = LRUCache(max_memory_bytes=max_cache_size)
+    return _result_cache
+
+
+def clear_cache() -> None:
+    """Clear the global result cache.
+
+    Example:
+        >>> clear_cache()
+        >>> # All cached results released
+    """
+    cache = get_result_cache()
+    cache.clear()
+
+
+def show_cache_stats() -> dict[str, int | float]:
+    """Show statistics for the global cache.
+
+    Returns:
+        Dictionary with cache statistics.
+
+    Example:
+        >>> stats = show_cache_stats()
+        >>> print(f"Hit rate: {stats['hit_rate']*100:.1f}%")
+        >>> print(f"Memory used: {stats['memory_mb']:.1f} MB")
+    """
+    cache = get_result_cache()
+    return cache.stats()
+
+
+def cache_key(*args: Any, **kwargs: Any) -> str:
+    """Generate cache key from arguments.
+
+    Args:
+        *args: Positional arguments.
+        **kwargs: Keyword arguments.
+
+    Returns:
+        Hash-based cache key.
+
+    Example:
+        >>> key = cache_key("fft", samples=1000, nfft=2048)
+        >>> # Use key for caching
+    """
+    # Create stable string representation
+    parts = [str(arg) for arg in args]
+    parts.extend(f"{k}={v}" for k, v in sorted(kwargs.items()))
+    key_str = "|".join(parts)
+
+    # Hash for consistent key
+    return hashlib.md5(key_str.encode()).hexdigest()
+
+
+# =============================================================================
+# HDF5 Lazy Loading (MEM-017)
+# =============================================================================
+
+
+def load_hdf5_lazy(
+    file_path: str,
+    dataset_path: str = "/data",
+) -> Any:
+    """Load HDF5 dataset as lazy h5py.Dataset (not fully in memory).
+
+    Enables partial loading via slicing without loading entire dataset.
+
+    Args:
+        file_path: Path to HDF5 file.
+        dataset_path: Path to dataset within file (default "/data").
+
+    Returns:
+        h5py.Dataset object (lazy, not loaded until accessed).
+
+    Raises:
+        ImportError: If h5py is not available.
+        FileNotFoundError: If file does not exist.
+        KeyError: If dataset not found in file.
+
+    Example:
+        >>> # Load dataset lazily
+        >>> dataset = load_hdf5_lazy("large_file.h5", "/signals/ch1")
+        >>> # Only load specific range (not entire file)
+        >>> chunk = dataset[1000:2000]
+        >>> print(f"Chunk shape: {chunk.shape}")
+
+    References:
+        MEM-017: HDF5 Chunked Dataset Access
+    """
+    try:
+        import h5py
+    except ImportError:
+        raise ImportError(  # noqa: B904
+            "h5py required for lazy HDF5 loading. Install with: pip install h5py"
+        )
+
+    from pathlib import Path
+
+    path = Path(file_path)
+    if not path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    # Open file in read mode
+    # Note: File handle should be kept open for lazy access
+    # User is responsible for closing the file handle
+    f = h5py.File(file_path, "r")
+
+    if dataset_path not in f:
+        available = list(f.keys())
+        f.close()
+        raise KeyError(
+            f"Dataset '{dataset_path}' not found in HDF5 file. "
+            f"Available datasets: {', '.join(available)}"
+        )
+
+    dataset = f[dataset_path]
+
+    return dataset
+
+
+class LazyHDF5Array:
+    """Wrapper for lazy HDF5 dataset access with context management.
+
+    Provides automatic file handle cleanup and numpy-like slicing.
+
+    Args:
+        file_path: Path to HDF5 file.
+        dataset_path: Path to dataset within file.
+
+    Example:
+        >>> with LazyHDF5Array("data.h5", "/signals/ch1") as arr:
+        ...     # Only loads specific slice
+        ...     chunk = arr[1000:2000]
+        ...     print(f"Shape: {arr.shape}, dtype: {arr.dtype}")
+        >>> # File automatically closed
+
+    References:
+        MEM-017: HDF5 Chunked Dataset Access
+        MEM-019: Explicit Resource Cleanup
+    """
+
+    def __init__(self, file_path: str, dataset_path: str = "/data"):
+        self._file_path = file_path
+        self._dataset_path = dataset_path
+        self._file = None
+        self._dataset = None
+
+    def __enter__(self) -> LazyHDF5Array:
+        """Open HDF5 file and dataset."""
+        try:
+            import h5py
+        except ImportError:
+            raise ImportError("h5py required. Install with: pip install h5py")  # noqa: B904
+
+        self._file = h5py.File(self._file_path, "r")
+        self._dataset = self._file[self._dataset_path]  # type: ignore[index]
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Close HDF5 file."""
+        # Note: exc_val and exc_tb intentionally unused but required for Python 3.11+ compatibility
+        if self._file is not None:
+            self._file.close()  # type: ignore[unreachable]
+            self._file = None
+            self._dataset = None
+
+    def __getitem__(self, key: Any) -> NDArray[Any]:
+        """Get item/slice from dataset (triggers partial load)."""
+        if self._dataset is None:
+            raise RuntimeError("LazyHDF5Array must be used as context manager")
+        return np.asarray(self._dataset[key])  # type: ignore[unreachable]
+
+    @property
+    def shape(self) -> tuple[int, ...]:
+        """Dataset shape."""
+        if self._dataset is None:
+            raise RuntimeError("LazyHDF5Array must be used as context manager")
+        return self._dataset.shape  # type: ignore[unreachable]
+
+    @property
+    def dtype(self) -> np.dtype[Any]:
+        """Dataset dtype."""
+        if self._dataset is None:
+            raise RuntimeError("LazyHDF5Array must be used as context manager")
+        return self._dataset.dtype  # type: ignore[unreachable]
+
+    @property
+    def size(self) -> int:
+        """Total number of elements."""
+        if self._dataset is None:
+            raise RuntimeError("LazyHDF5Array must be used as context manager")
+        return self._dataset.size  # type: ignore[unreachable]
+
+    def __len__(self) -> int:
+        """Length of first dimension."""
+        if self._dataset is None:
+            raise RuntimeError("LazyHDF5Array must be used as context manager")
+        return len(self._dataset)  # type: ignore[unreachable]
+
+
+__all__ = [
+    "ArrayManager",
+    "LRUCache",
+    "LazyHDF5Array",
+    "ResourceManager",
+    "cache_key",
+    "clear_cache",
+    "get_result_cache",
+    "load_hdf5_lazy",
+    "show_cache_stats",
+]