nshutils 0.22.6__tar.gz → 0.31.0__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "nshutils"
-version = "0.22.6"
+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.22.6 → 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.22.6 → 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.22.6 → nshutils-0.31.0}/src/nshutils/actsave/_saver.py

@@ -241,6 +241,17 @@ class _Saver:
 class ActSaveProvider:
     _saver: _Saver | None = None
     _prefixes: list[str] = []
+    _disable_count: int = 0
+
+    @property
+    def is_initialized(self) -> bool:
+        """Returns True if ActSave.enable() has been called and not subsequently disabled."""
+        return self._saver is not None
+
+    @property
+    def is_enabled(self) -> bool:
+        """Returns True if ActSave is currently active and will save activations."""
+        return self.is_initialized and self._disable_count == 0
 
     def enable(self, save_dir: Path | None = None):
         """
@@ -294,6 +305,34 @@ class ActSaveProvider:
 
         self._saver = None
         self._prefixes = []
+        self._disable_count = 0
+
+    @contextlib.contextmanager
+    def disabled(self, condition: bool | Callable[[], bool] = True):
+        """
+        Context manager to temporarily disable activation saving.
+
+        Args:
+            condition (bool | Callable[[], bool], optional):
+                If True or a callable returning True, saving is disabled within this context.
+                Defaults to True.
+        """
+        if _torch_is_scripting():
+            yield
+            return
+
+        should_disable = condition() if callable(condition) else condition
+        if should_disable:
+            self._disable_count += 1
+
+        try:
+            yield
+        finally:
+            if should_disable:
+                self._disable_count -= 1
+                if self._disable_count < 0:  # Should not happen
+                    log.warning("ActSave disable count went below zero.")
+                    self._disable_count = 0
 
     @contextlib.contextmanager
     def context(self, label: str):
@@ -307,7 +346,7 @@ class ActSaveProvider:
             yield
             return
 
-        if self._saver is None:
+        if not self.is_enabled:
             yield
             return
 
@@ -368,9 +407,12 @@ class ActSaveProvider:
         if _torch_is_scripting():
             return
 
-        if self._saver is None:
+        if not self.is_enabled:
             return
 
+        # Ensure _saver is not None, which is guaranteed by is_enabled but mypy needs help
+        assert self._saver is not None
+
         if acts is not None and callable(acts):
             acts = acts()
         self._saver.save(acts, **kwargs)

{nshutils-0.22.6 → 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.31.0/src/nshutils/lovely/_base.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import contextlib
+import functools
+import importlib.util
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, Generic, Optional, cast
+
+from typing_extensions import (
+    ParamSpec,
+    Protocol,
+    TypeAliasType,
+    TypeVar,
+    override,
+    runtime_checkable,
+)
+
+from .utils import LovelyStats, format_tensor_stats
+
+log = logging.getLogger(__name__)
+
+TArray = TypeVar("TArray", infer_variance=True)
+P = ParamSpec("P")
+
+LovelyStatsFn = TypeAliasType(
+    "LovelyStatsFn",
+    Callable[[TArray], Optional[LovelyStats]],
+    type_params=(TArray,),
+)
+
+
+@runtime_checkable
+class LovelyReprFn(Protocol[TArray]):
+    @property
+    def __lovely_repr_instance__(self) -> lovely_repr[TArray]: ...
+
+    @__lovely_repr_instance__.setter
+    def __lovely_repr_instance__(self, value: lovely_repr[TArray]) -> None: ...
+
+    @property
+    def __name__(self) -> str: ...
+
+    def set_fallback_repr(self, repr_fn: Callable[[TArray], str]) -> None: ...
+    def __call__(self, value: TArray, /) -> str: ...
+
+
+def _find_missing_deps(dependencies: list[str]):
+    missing_deps: list[str] = []
+
+    for dep in dependencies:
+        if importlib.util.find_spec(dep) is not None:
+            continue
+
+        missing_deps.append(dep)
+
+    return missing_deps
+
+
+class lovely_repr(Generic[TArray]):
+    @override
+    def __init__(
+        self,
+        dependencies: list[str],
+        fallback_repr: Callable[[TArray], str] | None = None,
+    ):
+        """
+        Decorator to create a lovely representation 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.
+            fallback_repr: A function that takes an array and returns its fallback representation.
+        Returns:
+            A decorator function that takes a function and returns a lovely representation function.
+
+        Example:
+            @lovely_repr(dependencies=["torch"])
+            def my_array_stats(array):
+                return {...}
+        """
+        super().__init__()
+
+        if fallback_repr is None:
+            fallback_repr = repr
+
+        self._dependencies = dependencies
+        self._fallback_repr = fallback_repr
+
+    def set_fallback_repr(self, repr_fn: Callable[[TArray], str]) -> None:
+        self._fallback_repr = repr_fn
+
+    def __call__(
+        self, array_stats_fn: LovelyStatsFn[TArray], /
+    ) -> LovelyReprFn[TArray]:
+        @functools.wraps(array_stats_fn)
+        def wrapper_fn(array: TArray) -> str:
+            if missing_deps := _find_missing_deps(self._dependencies):
+                log.warning(
+                    f"Missing dependencies: {', '.join(missing_deps)}. "
+                    "Skipping lovely representation."
+                )
+                return self._fallback_repr(array)
+
+            if (stats := array_stats_fn(array)) is None:
+                return self._fallback_repr(array)
+
+            return format_tensor_stats(stats)
+
+        wrapper = cast(LovelyReprFn[TArray], wrapper_fn)
+        wrapper.__lovely_repr_instance__ = self
+        wrapper.set_fallback_repr = self.set_fallback_repr
+        return wrapper
+
+
+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 []
+
+    @abstractmethod
+    def patch(self):
+        """Subclasses must implement this."""
+
+    @abstractmethod
+    def unpatch(self):
+        """Subclasses must implement this."""
+
+    @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 self
+
+        self.patch()
+        self._patched = True
+        return self
+
+    @override
+    def __exit__(self, *exc_info):
+        if not self._patched:
+            return
+
+        self.unpatch()
+        self._patched = False
+
+    def close(self):
+        """Explicitly clean up the resource."""
+        self.__exit__(None, None, None)

{nshutils-0.22.6 → 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.22.6 → 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:
@@ -53,7 +54,7 @@ def _device(array: jax.Array) -> str:
     return f"{device.platform}:{device.id}"
 
 
-@lovely_repr(dependencies=["jax"], fallback_repr=jax.Array.__repr__)
+@lovely_repr(dependencies=["jax"])
 def jax_repr(array: jax.Array) -> LovelyStats | None:
     import jax.numpy as jnp
 
@@ -77,17 +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__
-    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.22.6 → 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
 
 
@@ -51,7 +52,7 @@ def _dtype_str(array: np.ndarray) -> str:
     return dtype_base
 
 
-@lovely_repr(dependencies=["numpy"], fallback_repr=np.array_repr)
+@lovely_repr(dependencies=["numpy"])
 def numpy_repr(array: np.ndarray) -> LovelyStats | None:
     # For dtypes like `object` or `str`, we let the fallback repr handle it
     if not np.issubdtype(array.dtype, np.number):
@@ -71,42 +72,44 @@ def numpy_repr(array: np.ndarray) -> LovelyStats | None:
     }
 
 
-# If numpy 2.0, use the new API override_repr.
-if _np_ge_2():
+numpy_repr.set_fallback_repr(np.array_repr)
 
-    @monkey_patch_contextmanager(dependencies=["numpy"])
-    def numpy_monkey_patch():
-        try:
+
+class numpy_monkey_patch(lovely_patch):
+    @override
+    def dependencies(self) -> list[str]:
+        return ["numpy"]
+
+    @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.22.6 → 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:
@@ -59,7 +60,7 @@ def _to_np(tensor: torch.Tensor) -> np.ndarray:
     return t_np
 
 
-@lovely_repr(dependencies=["torch"], fallback_repr=torch.Tensor.__repr__)
+@lovely_repr(dependencies=["torch"])
 def torch_repr(tensor: torch.Tensor) -> LovelyStats | None:
     return {
         # Basic attributes
@@ -80,20 +81,31 @@ 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
+
+        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)
 
-    original_repr = torch.Tensor.__repr__
-    original_str = torch.Tensor.__str__
-    original_parameter_repr = torch.nn.Parameter.__repr__
-    try:
         patch_to(torch.Tensor, "__repr__", torch_repr)
         patch_to(torch.Tensor, "__str__", torch_repr)
-        del torch.nn.Parameter.__repr__
-
-        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)
+        try:
+            delattr(torch.nn.Parameter, "__repr__")
+        except AttributeError:
+            pass
+
+    @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.22.6/src/nshutils/lovely/_base.py

@@ -1,158 +0,0 @@
-from __future__ import annotations
-
-import functools
-import importlib.util
-import logging
-from collections.abc import Callable, Iterator
-
-from typing_extensions import ParamSpec, TypeAliasType, TypeVar
-
-from ..util import ContextResource, resource_factory_contextmanager
-from .utils import LovelyStats, format_tensor_stats
-
-log = logging.getLogger(__name__)
-
-TArray = TypeVar("TArray", infer_variance=True)
-P = ParamSpec("P")
-
-LovelyStatsFn = TypeAliasType(
-    "LovelyStatsFn",
-    Callable[[TArray], LovelyStats | None],
-    type_params=(TArray,),
-)
-LovelyReprFn = TypeAliasType(
-    "LovelyReprFn",
-    Callable[[TArray], str],
-    type_params=(TArray,),
-)
-
-
-def _find_missing_deps(dependencies: list[str]):
-    missing_deps: list[str] = []
-
-    for dep in dependencies:
-        if importlib.util.find_spec(dep) is not None:
-            continue
-
-        missing_deps.append(dep)
-
-    return missing_deps
-
-
-def lovely_repr(dependencies: list[str], fallback_repr: Callable[[TArray], str]):
-    """
-    Decorator to create a lovely representation 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.
-        fallback_repr: A function that takes an array and returns its fallback representation.
-    Returns:
-        A decorator function that takes a function and returns a lovely representation function.
-
-    Example:
-        @lovely_repr(dependencies=["torch"])
-        def my_array_stats(array):
-            return {...}
-    """
-
-    def decorator_fn(array_stats_fn: LovelyStatsFn[TArray]) -> LovelyReprFn[TArray]:
-        """
-        Decorator to create a lovely representation function for an array.
-
-        Args:
-            array_stats_fn: A function that takes an array and returns its stats,
-                or `None` if the array is not supported.
-
-        Returns:
-            A function that takes an array and returns its lovely representation.
-        """
-
-        @functools.wraps(array_stats_fn)
-        def wrapper(array: TArray) -> str:
-            if missing_deps := _find_missing_deps(dependencies):
-                log.warning(
-                    f"Missing dependencies: {', '.join(missing_deps)}. "
-                    "Skipping lovely representation."
-                )
-                return fallback_repr(array)
-
-            if (stats := array_stats_fn(array)) is None:
-                return fallback_repr(array)
-
-            return format_tensor_stats(stats)
-
-        return wrapper
-
-    return decorator_fn
-
-
-LovelyMonkeyPatchInputFn = TypeAliasType(
-    "LovelyMonkeyPatchInputFn",
-    Callable[P, Iterator[None]],
-    type_params=(P,),
-)
-LovelyMonkeyPatchFn = TypeAliasType(
-    "LovelyMonkeyPatchFn",
-    Callable[P, ContextResource[None]],
-    type_params=(P,),
-)
-
-
-def _nullcontext_generator():
-    """A generator that does nothing."""
-    yield
-
-
-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):
-            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.
-
-    Returns:
-        A decorator function that takes a function and returns a monkey patch function.
-
-    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.
-        """
-
-        wrapped_fn = _wrap_monkey_patch_fn(monkey_patch_fn, dependencies)
-        return resource_factory_contextmanager(wrapped_fn)
-
-    return decorator_fn

nshutils-0.22.6/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))