nshutils 0.17.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

nshutils/__init__.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 from . import actsave as actsave
 from . import typecheck as typecheck
-from ._display import display as display
 from .actsave import ActLoad as ActLoad
 from .actsave import ActSave as ActSave
+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 lovely as lovely
 from .logging import pretty as pretty

nshutils/actsave/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from ._loader import ActLoad as ActLoad
 from ._saver import Activation as Activation
 from ._saver import ActSave as ActSave

nshutils/actsave/_loader.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import pprint
 from dataclasses import dataclass, field
 from functools import cached_property

nshutils/actsave/_saver.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import fnmatch
 import tempfile

nshutils/collections.py

@@ -2,6 +2,8 @@
 # Licensed under the Apache License, Version 2.0 (the "License");
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
+from __future__ import annotations
+
 import dataclasses
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable, Mapping, Sequence

nshutils/{_display.py → display.py}

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import importlib.util
 from functools import cache
 from typing import Any
@@ -25,7 +27,7 @@ def _rich_installed():
 
 def display(*args: Any):
     """
-        Display the given arguments in the current environment.
+    Display the given arguments in the current environment.
 
     If executed in an IPython environment, the display will be handled
     by treescope if installed, or rich if available. If neither are

nshutils/logging.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from pathlib import Path
 
@@ -7,7 +9,7 @@ def init_python_logging(
     lovely_tensors: bool = False,
     lovely_numpy: bool = False,
     treescope: bool = False,
-    treescope_autovisualize_arrays: bool = True,
+    treescope_autovisualize_arrays: bool = False,
     rich: bool = False,
     rich_tracebacks: bool = False,
     log_level: int | str | None = logging.INFO,
@@ -81,8 +83,8 @@ def pretty(
     *,
     lovely_tensors: bool = True,
     lovely_numpy: bool = True,
-    treescope: bool = True,
-    treescope_autovisualize_arrays: bool = True,
+    treescope: bool = False,
+    treescope_autovisualize_arrays: bool = False,
     log_level: int | str | None = logging.INFO,
     log_save_dir: Path | None = None,
     rich_log_handler: bool = False,
@@ -104,8 +106,8 @@ def lovely(
     *,
     lovely_tensors: bool = True,
     lovely_numpy: bool = True,
-    treescope: bool = True,
-    treescope_autovisualize_arrays: bool = True,
+    treescope: bool = False,
+    treescope_autovisualize_arrays: bool = False,
     log_level: int | str | None = logging.INFO,
     log_save_dir: Path | None = None,
     rich_log_handler: bool = False,

nshutils/snoop.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import contextlib
 import importlib.metadata
 from typing import Any, Protocol, cast
@@ -13,6 +15,7 @@ class SnoopConstructor(Protocol):
     def disable(self) -> contextlib.AbstractContextManager: ...
 
 
+snoop: SnoopConstructor
 try:
     import warnings
     from contextlib import nullcontext
@@ -201,7 +204,7 @@ try:
         disable = nullcontext
         __call__ = TorchSnooper
 
-    snoop: SnoopConstructor = cast(Any, _Snoop())
+    snoop = cast(Any, _Snoop())
 
 except ImportError:
     import warnings
@@ -227,4 +230,4 @@ except ImportError:
 
             return super().__enter__()
 
-    snoop: SnoopConstructor = cast(Any, _snoop_cls)
+    snoop = cast(Any, _snoop_cls)

nshutils/typecheck.py

@@ -1,6 +1,10 @@
+from __future__ import annotations
+
 import os
+import sys
 from collections.abc import Sequence
 from logging import getLogger
+from types import FrameType as _FrameType
 from typing import Any
 
 from jaxtyping import BFloat16 as BFloat16
@@ -76,6 +80,21 @@ def typecheck_modules(modules: Sequence[str]):
     log.critical(f"Type checking the following modules: {modules}")
 
 
+def _get_frame_package_name_or_none(frame: _FrameType) -> str | None:
+    # Taken from `beartype._util.func.utilfuncframe.get_frame_package_name_or_none`.
+    assert isinstance(frame, _FrameType), f"{repr(frame)} not stack frame."
+
+    # Fully-qualified name of the parent package of the child module declaring
+    # the callable whose code object is that of this stack frame's if that
+    # module declares its name *OR* the empty string otherwise (e.g., if that
+    # module is either a top-level module or script residing outside any parent
+    # package structure).
+    frame_package_name = frame.f_globals.get("__package__")
+
+    # Return the name of this parent package.
+    return frame_package_name
+
+
 def typecheck_this_module(additional_modules: Sequence[str] = ()):
     """
     Typecheck the calling module and any additional modules using `jaxtyping`.
@@ -86,13 +105,11 @@ def typecheck_this_module(additional_modules: Sequence[str] = ()):
     # Get the calling module's name.
     # Here, we can just use beartype's internal implementation behind
     # `beartype_this_package`.
-    from beartype._util.func.utilfuncframe import get_frame, get_frame_package_name
 
     # Get the calling module's name.
-    assert get_frame is not None, "get_frame is None"
-    frame = get_frame(1)
+    frame = sys._getframe(1)
     assert frame is not None, "frame is None"
-    calling_module_name = get_frame_package_name(frame)
+    calling_module_name = _get_frame_package_name_or_none(frame)
     assert calling_module_name is not None, "calling_module_name is None"
 
     # Typecheck the calling module + any additional modules.

{nshutils-0.17.0.dist-info → nshutils-0.19.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: nshutils
-Version: 0.17.0
+Version: 0.19.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -9,6 +9,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Provides-Extra: extra
 Requires-Dist: beartype
 Requires-Dist: jaxtyping
@@ -20,7 +21,7 @@ Requires-Dist: rich ; extra == "extra"
 Requires-Dist: treescope ; extra == "extra"
 Requires-Dist: typing-extensions
 Requires-Dist: uuid7
-Project-URL: homepage, https://github.com/nimashoghi/nshutils
+Project-URL: Homepage, https://github.com/nimashoghi/nshutils
 Description-Content-Type: text/markdown
 
 # nshutils

nshutils-0.19.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+nshutils/__init__.py,sha256=4sIp11UG2TFbvfRKbWdVUX4lzIKlsYEJuFK34735rSQ,642
+nshutils/actsave/__init__.py,sha256=hAVsog9d1g3_rQN1TRslrl6sK1PhCGbjy8PPUAmJI58,203
+nshutils/actsave/_loader.py,sha256=mof3HezUNvLliz7macstX6ewXW05L0Mtv3zJyrbmImg,4640
+nshutils/actsave/_saver.py,sha256=Kor7PEk__noRDEAfCZ_I4vxql5UfWSIQIQ1OmW2RRTI,10290
+nshutils/collections.py,sha256=QWGyANmo4Efq4XRNHDSTE9tRLStwEZHGwE0ATHR-Vqo,5233
+nshutils/display.py,sha256=Ge63yllx7gi-MKL3mKQeQ5doql_nj56-o5aoTVmusDg,1473
+nshutils/logging.py,sha256=-6IB0GTDDS8ue1H2tzkv_OLf4bZVN1ywL08TlDZWbtQ,3737
+nshutils/snoop.py,sha256=7d7_Q5sJmINL1J29wcnxEvpV95zvZYNoVn5frCq-rww,7393
+nshutils/typecheck.py,sha256=UOUYfa72wTmc-a7VQw52tKFb4U10xq1qcZuEWc2sAd8,5588
+nshutils-0.19.0.dist-info/METADATA,sha256=g7LvXvnncRKS7MJJ6BmZHZ4aMHVbb6X77jJx6wD9Dwk,4168
+nshutils-0.19.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshutils-0.19.0.dist-info/RECORD,,

{nshutils-0.17.0.dist-info → nshutils-0.19.0.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.9.0
+Generator: poetry-core 2.1.1
 Root-Is-Purelib: true
 Tag: py3-none-any

nshutils-0.17.0.dist-info/RECORD

@@ -1,12 +0,0 @@
-nshutils/__init__.py,sha256=T1MnsjctEI1c5wM06d_hUOuGIYLIz_NEb1e_tT-TgQ0,539
-nshutils/_display.py,sha256=6CSaKdsFBHBGQXMppumRfO5fOfiJDXQ7FhtAvN8id7M,1441
-nshutils/actsave/__init__.py,sha256=6gJ49011Ad3kS8BejeZRPj7ZyVyIcmX-VKLEmYZyGM8,167
-nshutils/actsave/_loader.py,sha256=fAhD32DrJa4onkYfcwc21YIeGEYzOSXCK_HVo9SZLgQ,4604
-nshutils/actsave/_saver.py,sha256=-uKMmMKjEErCa8pfLhAW4077GFcUjxr0Qq0PFAnbEOw,10254
-nshutils/collections.py,sha256=EE_qLd-LrsX5lsyk9GSKh03Q8bhn9CHB3jiEeNj4uF4,5197
-nshutils/logging.py,sha256=n5v2sNmSj5gKkdjL4avP4OpoUzJb4BVtdDtTm_dzVgM,3696
-nshutils/snoop.py,sha256=JW2JZg6oZd9WWubOP6ok6GUbh5mIQUOdP_T29iH5RCE,7369
-nshutils/typecheck.py,sha256=_KtfinRy9A0Dgq78kN5MGGrefvb6jn2tGY6svdLEBAs,4915
-nshutils-0.17.0.dist-info/METADATA,sha256=LF0Opz5cgFBshbBKuAbfAZJd011ceKdA2_sDvCA2qzA,4117
-nshutils-0.17.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshutils-0.17.0.dist-info/RECORD,,