nshutils 0.30.1__tar.gz → 0.31.0__tar.gz

{nshutils-0.30.1 → nshutils-0.31.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshutils
-Version: 0.30.1
+Version: 0.31.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshutils-0.30.1 → nshutils-0.31.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshutils"
-version = "0.30.1"
+version = "0.31.0"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.9,<4.0"
@@ -28,6 +28,8 @@ basedpyright = "*"
 ruff = "*"
 ipykernel = "*"
 ipywidgets = "*"
+pytest = "*"
+pytest-cov = "*"
 
 [build-system]
 requires = ["poetry-core"]
@@ -47,3 +49,9 @@ ignore = ["F722", "F821", "E731", "E741"]
 
 [tool.ruff.lint.isort]
 required-imports = ["from __future__ import annotations"]
+
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+addopts = "--cov=nshutils --cov-report=term-missing"

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/__init__.pyi

@@ -7,6 +7,7 @@ from .collections import apply_to_collection as apply_to_collection
 from .display import display as display
 from .logging import init_python_logging as init_python_logging
 from .logging import setup_logging as setup_logging
+from .lovely import lovely_monkey_patch as lovely_monkey_patch
 from .snoop import snoop as snoop
 from .typecheck import tassert as tassert
 from .typecheck import typecheck_modules as typecheck_modules

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/actsave/_loader.py

@@ -105,14 +105,41 @@ class ActLoad:
         path, _ = max(all_versions, key=lambda p: p[1])
         return cls(path)
 
-    def __init__(self, dir: Path):
+    def __init__(
+        self,
+        dir: Path | None = None,
+        *,
+        _base_activations: dict[str, LoadedActivation] | None = None,
+        _prefix_chain: list[str] | None = None,
+    ):
+        """Initialize ActLoad from a directory or from filtered activations.
+
+        Args:
+            dir: Path to the activation directory. Required for root ActLoad instances.
+            _base_activations: Pre-filtered activations dict. Used internally for prefix filtering.
+            _prefix_chain: Chain of prefixes that have been applied. Used for repr.
+        """
         self._dir = dir
+        self._base_activations = _base_activations
+        self._prefix_chain = _prefix_chain or []
 
     def activation(self, name: str):
-        return LoadedActivation(self._dir, name)
+        if self._dir is None:
+            raise ValueError(
+                "Cannot create activation from filtered ActLoad without base directory"
+            )
+        # For filtered instances, we need to reconstruct the full name
+        full_name = "".join(self._prefix_chain) + name
+        return LoadedActivation(self._dir, full_name)
 
     @cached_property
     def activations(self):
+        if self._base_activations is not None:
+            return self._base_activations
+
+        if self._dir is None:
+            raise ValueError("ActLoad requires either dir or _base_activations")
+
         dirs = list(self._dir.iterdir())
         # Sort the dirs by the last modified time
         dirs.sort(key=lambda p: p.stat().st_mtime)
@@ -128,6 +155,9 @@ class ActLoad:
     def __len__(self):
         return len(self.activations)
 
+    def _ipython_key_completions_(self):
+        return list(self.activations.keys())
+
     @override
     def __repr__(self):
         acts_str = pprint.pformat(
@@ -137,10 +167,50 @@ class ActLoad:
             }
         )
         acts_str = acts_str.replace("'<", "<").replace(">'", ">")
-        return f"ActLoad({acts_str})"
+
+        if self._prefix_chain:
+            prefix_str = "".join(self._prefix_chain)
+            return f"ActLoad(prefix='{prefix_str}', {acts_str})"
+        else:
+            return f"ActLoad({acts_str})"
 
     def get(self, name: str, /, default: T) -> LoadedActivation | T:
         return self.activations.get(name, default)
 
+    def filter_by_prefix(self, prefix: str) -> ActLoad:
+        """Create a filtered view of activations that match the given prefix.
+
+        Args:
+            prefix: The prefix to filter by. Only activations whose names start
+                   with this prefix will be included in the filtered view.
+
+        Returns:
+            A new ActLoad instance that provides access to matching activations
+            with the prefix stripped from their keys. Can be chained multiple times.
+
+        Example:
+            >>> loader = ActLoad(some_dir)
+            >>> # If loader has keys "my.activation.first", "my.activation.second", "other.key"
+            >>> filtered = loader.filter_by_prefix("my.activation.")
+            >>> filtered["first"]  # Accesses "my.activation.first"
+            >>> filtered["second"]  # Accesses "my.activation.second"
+            >>> # Can be chained:
+            >>> double_filtered = loader.filter_by_prefix("my.").filter_by_prefix("activation.")
+            >>> double_filtered["first"]  # Also accesses "my.activation.first"
+        """
+        filtered_activations = {}
+        for name, activation in self.activations.items():
+            if name.startswith(prefix):
+                # Strip the prefix from the key
+                stripped_name = name[len(prefix) :]
+                filtered_activations[stripped_name] = activation
+
+        new_prefix_chain = self._prefix_chain + [prefix]
+        return ActLoad(
+            _base_activations=filtered_activations,
+            _prefix_chain=new_prefix_chain,
+            dir=self._dir,
+        )
+
 
 ActivationLoader = ActLoad

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/lovely/__init__.py

@@ -8,3 +8,5 @@ from .numpy_ import numpy_monkey_patch as numpy_monkey_patch
 from .numpy_ import numpy_repr as numpy_repr
 from .torch_ import torch_monkey_patch as torch_monkey_patch
 from .torch_ import torch_repr as torch_repr
+
+lovely_monkey_patch = monkey_patch

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/lovely/_base.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import contextlib
 import functools
 import importlib.util
 import logging
-from collections.abc import Callable, Iterator
-from typing import Generic, Optional, cast
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Generic, Optional, cast
 
 from typing_extensions import (
     ParamSpec,
@@ -15,7 +17,6 @@ from typing_extensions import (
     runtime_checkable,
 )
 
-from ..util import ContextResource, resource_factory_contextmanager
 from .utils import LovelyStats, format_tensor_stats
 
 log = logging.getLogger(__name__)
@@ -113,72 +114,47 @@ class lovely_repr(Generic[TArray]):
         return wrapper
 
 
-LovelyMonkeyPatchInputFn = TypeAliasType(
-    "LovelyMonkeyPatchInputFn",
-    Callable[P, Iterator[None]],
-    type_params=(P,),
-)
-LovelyMonkeyPatchFn = TypeAliasType(
-    "LovelyMonkeyPatchFn",
-    Callable[P, ContextResource[None]],
-    type_params=(P,),
-)
+class lovely_patch(contextlib.AbstractContextManager["lovely_patch"], ABC):
+    def __init__(self):
+        self._patched = False
+        self.__enter__()
 
+    def dependencies(self) -> list[str]:
+        """Subclasses can override this to specify the dependencies of the patch."""
+        return []
 
-def _nullcontext_generator():
-    """A generator that does nothing."""
-    yield
+    @abstractmethod
+    def patch(self):
+        """Subclasses must implement this."""
 
+    @abstractmethod
+    def unpatch(self):
+        """Subclasses must implement this."""
 
-def _wrap_monkey_patch_fn(
-    monkey_patch_fn: LovelyMonkeyPatchInputFn[P],
-    dependencies: list[str],
-) -> LovelyMonkeyPatchInputFn[P]:
-    @functools.wraps(monkey_patch_fn)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> Iterator[None]:
-        if missing_deps := _find_missing_deps(dependencies):
+    @override
+    def __enter__(self):
+        if self._patched:
+            return self
+
+        if missing_deps := _find_missing_deps(self.dependencies()):
             log.warning(
                 f"Missing dependencies: {', '.join(missing_deps)}. "
                 "Skipping monkey patch."
             )
-            return _nullcontext_generator()
-
-        return monkey_patch_fn(*args, **kwargs)
-
-    return wrapper
-
-
-def monkey_patch_contextmanager(dependencies: list[str]):
-    """
-    Decorator to create a monkey patch function for an array.
-
-    Args:
-        dependencies: List of dependencies to check before running the function.
-            If any dependency is not available, the function will not run.
+            return self
 
-    Returns:
-        A decorator function that takes a function and returns a monkey patch function.
+        self.patch()
+        self._patched = True
+        return self
 
-    Example:
-        @monkey_patch_contextmanager(dependencies=["torch"])
-        def my_array_monkey_patch():
-            ...
-    """
-
-    def decorator_fn(
-        monkey_patch_fn: LovelyMonkeyPatchInputFn[P],
-    ) -> LovelyMonkeyPatchFn[P]:
-        """
-        Decorator to create a monkey patch function for an array.
-
-        Args:
-            monkey_patch_fn: A function that applies the monkey patch.
-
-        Returns:
-            A function that applies the monkey patch.
-        """
+    @override
+    def __exit__(self, *exc_info):
+        if not self._patched:
+            return
 
-        wrapped_fn = _wrap_monkey_patch_fn(monkey_patch_fn, dependencies)
-        return resource_factory_contextmanager(wrapped_fn)
+        self.unpatch()
+        self._patched = False
 
-    return decorator_fn
+    def close(self):
+        """Explicitly clean up the resource."""
+        self.__exit__(None, None, None)

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/lovely/_monkey_patch_all.py

@@ -5,9 +5,9 @@ import importlib.util
 import logging
 from typing import Literal
 
-from typing_extensions import TypeAliasType, assert_never
+from typing_extensions import TypeAliasType, assert_never, override
 
-from ..util import resource_factory_contextmanager
+from ._base import lovely_patch
 
 Library = TypeAliasType("Library", Literal["numpy", "torch", "jax"])
 
@@ -28,40 +28,48 @@ def _find_deps() -> list[Library]:
     return deps
 
 
-@resource_factory_contextmanager
-def monkey_patch(libraries: list[Library] | Literal["auto"] = "auto"):
-    if libraries == "auto":
-        libraries = _find_deps()
+class monkey_patch(lovely_patch):
+    def __init__(self, libraries: list[Library] | Literal["auto"] = "auto"):
+        self.libraries = libraries
+        if self.libraries == "auto":
+            self.libraries = _find_deps()
 
-    if not libraries:
-        raise ValueError(
-            "No libraries found for monkey patching. "
-            "Please install numpy, torch, or jax."
-        )
+        if not self.libraries:
+            raise ValueError(
+                "No libraries found for monkey patching. "
+                "Please install numpy, torch, or jax."
+            )
+
+        self.stack = contextlib.ExitStack()
+        super().__init__()
 
-    with contextlib.ExitStack() as stack:
-        for library in libraries:
+    @override
+    def patch(self):
+        for library in self.libraries:
             if library == "torch":
                 from .torch_ import torch_monkey_patch
 
-                stack.enter_context(torch_monkey_patch())
+                self.stack.enter_context(torch_monkey_patch())
             elif library == "jax":
                 from .jax_ import jax_monkey_patch
 
-                stack.enter_context(jax_monkey_patch())
+                self.stack.enter_context(jax_monkey_patch())
             elif library == "numpy":
                 from .numpy_ import numpy_monkey_patch
 
-                stack.enter_context(numpy_monkey_patch())
+                self.stack.enter_context(numpy_monkey_patch())
             else:
-                assert_never(library)
+                assert_never(library)  # type: ignore
 
         log.info(
-            f"Monkey patched libraries: {', '.join(libraries)}. "
+            f"Monkey patched libraries: {', '.join(self.libraries)}. "
             "You can now use the lovely functions with these libraries."
         )
-        yield
+
+    @override
+    def unpatch(self):
+        self.stack.close()
         log.info(
-            f"Unmonkey patched libraries: {', '.join(libraries)}. "
+            f"Unmonkey patched libraries: {', '.join(self.libraries)}. "
             "You can now use the lovely functions with these libraries."
         )

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/lovely/jax_.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, cast
 
 import numpy as np
+from typing_extensions import override
 
-from ._base import lovely_repr, monkey_patch_contextmanager
+from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats, patch_to
 
 if TYPE_CHECKING:
@@ -77,18 +78,25 @@ def jax_repr(array: jax.Array) -> LovelyStats | None:
     }
 
 
-@monkey_patch_contextmanager(dependencies=["jax"])
-def jax_monkey_patch():
-    from jax._src import array
+class jax_monkey_patch(lovely_patch):
+    @override
+    def dependencies(self) -> list[str]:
+        return ["jax"]
+
+    @override
+    def patch(self):
+        from jax._src import array
+
+        self.prev_repr = array.ArrayImpl.__repr__
+        self.prev_str = array.ArrayImpl.__str__
+        jax_repr.set_fallback_repr(self.prev_repr)
 
-    prev_repr = array.ArrayImpl.__repr__
-    prev_str = array.ArrayImpl.__str__
-    jax_repr.set_fallback_repr(prev_repr)
-    try:
         patch_to(array.ArrayImpl, "__repr__", jax_repr)
         patch_to(array.ArrayImpl, "__str__", jax_repr)
 
-        yield
-    finally:
-        patch_to(array.ArrayImpl, "__repr__", prev_repr)
-        patch_to(array.ArrayImpl, "__str__", prev_str)
+    @override
+    def unpatch(self):
+        from jax._src import array
+
+        patch_to(array.ArrayImpl, "__repr__", self.prev_repr)
+        patch_to(array.ArrayImpl, "__str__", self.prev_str)

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/lovely/numpy_.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 import logging
 
 import numpy as np
+from typing_extensions import override
 
-from ._base import lovely_repr, monkey_patch_contextmanager
+from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats
 
 
@@ -74,42 +75,41 @@ def numpy_repr(array: np.ndarray) -> LovelyStats | None:
 numpy_repr.set_fallback_repr(np.array_repr)
 
 
-# If numpy 2.0, use the new API override_repr.
-if _np_ge_2():
+class numpy_monkey_patch(lovely_patch):
+    @override
+    def dependencies(self) -> list[str]:
+        return ["numpy"]
 
-    @monkey_patch_contextmanager(dependencies=["numpy"])
-    def numpy_monkey_patch():
-        try:
+    @override
+    def patch(self):
+        if _np_ge_2():
+            self.original_options = np.get_printoptions()
             np.set_printoptions(override_repr=numpy_repr)
             logging.info(
                 f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
                 f"{np.get_printoptions()=}"
             )
-            yield
-        finally:
-            np.set_printoptions(override_repr=None)
+        else:
+            # For legacy numpy, `set_string_function(None)` reverts to the default,
+            # so no state needs to be saved.
+            np.set_string_function(numpy_repr, True)  # pyright: ignore[reportAttributeAccessIssue]
+            np.set_string_function(numpy_repr, False)  # pyright: ignore[reportAttributeAccessIssue]
             logging.info(
-                f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
                 f"{np.get_printoptions()=}"
             )
 
-else:
-
-    @monkey_patch_contextmanager(dependencies=["numpy"])
-    def numpy_monkey_patch():
-        try:
-            np.set_string_function(numpy_repr, True)  # pyright: ignore[reportAttributeAccessIssue]
-            np.set_string_function(numpy_repr, False)  # pyright: ignore[reportAttributeAccessIssue]
-
+    @override
+    def unpatch(self):
+        if _np_ge_2():
+            np.set_printoptions(**self.original_options)
             logging.info(
-                f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
                 f"{np.get_printoptions()=}"
             )
-            yield
-        finally:
+        else:
             np.set_string_function(None, True)  # pyright: ignore[reportAttributeAccessIssue]
             np.set_string_function(None, False)  # pyright: ignore[reportAttributeAccessIssue]
-
             logging.info(
                 f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
                 f"{np.get_printoptions()=}"

{nshutils-0.30.1 → nshutils-0.31.0}/src/nshutils/lovely/torch_.py

@@ -3,8 +3,9 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 import numpy as np
+from typing_extensions import override
 
-from ._base import lovely_repr, monkey_patch_contextmanager
+from ._base import lovely_patch, lovely_repr
 from .utils import LovelyStats, array_stats, patch_to
 
 if TYPE_CHECKING:
@@ -80,16 +81,20 @@ def torch_repr(tensor: torch.Tensor) -> LovelyStats | None:
     }
 
 
-@monkey_patch_contextmanager(dependencies=["torch"])
-def torch_monkey_patch():
-    import torch
+class torch_monkey_patch(lovely_patch):
+    @override
+    def dependencies(self) -> list[str]:
+        return ["torch"]
+
+    @override
+    def patch(self):
+        import torch
 
-    original_repr = torch.Tensor.__repr__
-    original_str = torch.Tensor.__str__
-    original_parameter_repr = torch.nn.Parameter.__repr__
-    torch_repr.set_fallback_repr(original_repr)
+        self.original_repr = torch.Tensor.__repr__
+        self.original_str = torch.Tensor.__str__
+        self.original_parameter_repr = torch.nn.Parameter.__repr__
+        torch_repr.set_fallback_repr(self.original_repr)
 
-    try:
         patch_to(torch.Tensor, "__repr__", torch_repr)
         patch_to(torch.Tensor, "__str__", torch_repr)
         try:
@@ -97,8 +102,10 @@ def torch_monkey_patch():
         except AttributeError:
             pass
 
-        yield
-    finally:
-        patch_to(torch.Tensor, "__repr__", original_repr)
-        patch_to(torch.Tensor, "__str__", original_str)
-        patch_to(torch.nn.Parameter, "__repr__", original_parameter_repr)
+    @override
+    def unpatch(self):
+        import torch
+
+        patch_to(torch.Tensor, "__repr__", self.original_repr)
+        patch_to(torch.Tensor, "__str__", self.original_str)
+        patch_to(torch.nn.Parameter, "__repr__", self.original_parameter_repr)

nshutils-0.30.1/src/nshutils/util.py

@@ -1,92 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import functools
-from collections.abc import Callable, Iterator
-from typing import Any, Generic
-
-from typing_extensions import ParamSpec, TypeVar, override
-
-R = TypeVar("R")
-P = ParamSpec("P")
-
-
-class ContextResource(contextlib.AbstractContextManager[R], Generic[R]):
-    """A class that provides both direct access to a resource and context management."""
-
-    def __init__(self, resource: R, cleanup_func: Callable[[R], Any]):
-        self.resource = resource
-        self._cleanup_func = cleanup_func
-
-    @override
-    def __enter__(self) -> R:
-        """When used as a context manager, return the wrapped resource."""
-        return self.resource
-
-    @override
-    def __exit__(self, *exc_info) -> None:
-        """Clean up the resource when exiting the context."""
-        self._cleanup_func(self.resource)
-
-    def close(self) -> None:
-        """Explicitly clean up the resource."""
-        self._cleanup_func(self.resource)
-
-
-def resource_factory(
-    create_func: Callable[P, R], cleanup_func: Callable[[R], None]
-) -> Callable[P, ContextResource[R]]:
-    """
-    Create a factory function that returns a ContextResource.
-
-    Args:
-        create_func: Function that creates the resource
-        cleanup_func: Function that cleans up the resource
-
-    Returns:
-        A function that returns a ContextResource wrapping the created resource
-    """
-
-    @functools.wraps(create_func)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> ContextResource[R]:
-        resource = create_func(*args, **kwargs)
-        return ContextResource(resource, cleanup_func)
-
-    return wrapper
-
-
-def resource_factory_from_context_fn(
-    context_func: Callable[P, contextlib.AbstractContextManager[R]],
-) -> Callable[P, ContextResource[R]]:
-    """
-    Create a factory function that returns a ContextResource.
-
-    Args:
-        context_func: Function that creates the resource
-
-    Returns:
-        A function that returns a ContextResource wrapping the created resource
-    """
-
-    @functools.wraps(context_func)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> ContextResource[R]:
-        context = context_func(*args, **kwargs)
-        resource = context.__enter__()
-        return ContextResource(resource, lambda _: context.__exit__(None, None, None))
-
-    return wrapper
-
-
-def resource_factory_contextmanager(
-    context_func: Callable[P, Iterator[R]],
-) -> Callable[P, ContextResource[R]]:
-    """
-    Create a factory function that returns a ContextResource.
-
-    Args:
-        context_func: Generator function that creates the resource, yields it, and cleans up the resource when done.
-
-    Returns:
-        A function that returns a ContextResource wrapping the created resource
-    """
-    return resource_factory_from_context_fn(contextlib.contextmanager(context_func))