nshutils 0.12.2__py3-none-any.whl → 0.14.0__py3-none-any.whl

nshutils/logging.py

@@ -6,6 +6,8 @@ def init_python_logging(
     *,
     lovely_tensors: bool = False,
     lovely_numpy: bool = False,
+    treescope: bool = False,
+    treescope_autovisualize_arrays: bool = True,
     rich: bool = False,
     rich_tracebacks: bool = False,
     log_level: int | str | None = logging.INFO,
@@ -31,6 +33,18 @@ def init_python_logging(
                 "Failed to import `lovely_numpy`. Ignoring pretty numpy array formatting"
             )
 
+    if treescope:
+        try:
+            import treescope as _treescope  # type: ignore
+
+            _treescope.basic_interactive_setup(
+                autovisualize_arrays=treescope_autovisualize_arrays
+            )
+        except ImportError:
+            logging.warning(
+                "Failed to import `treescope`. Ignoring `treescope` registration"
+            )
+
     log_handlers: list[logging.Handler] = []
     if log_save_dir:
         log_file = log_save_dir / "logging.log"
@@ -59,14 +73,18 @@ def pretty(
     *,
     lovely_tensors: bool = True,
     lovely_numpy: bool = True,
+    treescope: bool = True,
+    treescope_autovisualize_arrays: bool = True,
     log_level: int | str | None = logging.INFO,
     log_save_dir: Path | None = None,
-    rich_log_handler: bool = True,
-    rich_tracebacks: bool = True,
+    rich_log_handler: bool = False,
+    rich_tracebacks: bool = False,
 ):
     init_python_logging(
         lovely_tensors=lovely_tensors,
         lovely_numpy=lovely_numpy,
+        treescope=treescope,
+        treescope_autovisualize_arrays=treescope_autovisualize_arrays,
         rich=rich_log_handler,
         log_level=log_level,
         log_save_dir=log_save_dir,
@@ -78,14 +96,18 @@ def lovely(
     *,
     lovely_tensors: bool = True,
     lovely_numpy: bool = True,
+    treescope: bool = True,
+    treescope_autovisualize_arrays: bool = True,
     log_level: int | str | None = logging.INFO,
     log_save_dir: Path | None = None,
-    rich_log_handler: bool = True,
-    rich_tracebacks: bool = True,
+    rich_log_handler: bool = False,
+    rich_tracebacks: bool = False,
 ):
     pretty(
         lovely_tensors=lovely_tensors,
         lovely_numpy=lovely_numpy,
+        treescope=treescope,
+        treescope_autovisualize_arrays=treescope_autovisualize_arrays,
         log_level=log_level,
         log_save_dir=log_save_dir,
         rich_log_handler=rich_log_handler,

nshutils/snoop.py

@@ -1,4 +1,5 @@
 import contextlib
+import importlib.metadata
 from typing import Any, Protocol, cast
 
 from typing_extensions import TypeVar
@@ -18,7 +19,6 @@ try:
 
     import pysnooper  # type: ignore
     import pysnooper.utils  # type: ignore
-    from pkg_resources import DistributionNotFound, get_distribution
 
     try:
         import torch  # type: ignore
@@ -38,8 +38,8 @@ try:
         FLOATING_POINTS.add(getattr(torch, i))
 
     try:
-        __version__ = get_distribution(__name__).version
-    except DistributionNotFound:
+        __version__ = importlib.metadata.version(__name__)
+    except importlib.metadata.PackageNotFoundError:
         # package is not installed
         pass
 

{nshutils-0.12.2.dist-info → nshutils-0.14.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.12.2
+Version: 0.14.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -17,6 +17,7 @@ Requires-Dist: lovely-tensors ; extra == "extra"
 Requires-Dist: numpy
 Requires-Dist: pysnooper ; extra == "extra"
 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

{nshutils-0.12.2.dist-info → nshutils-0.14.0.dist-info}/RECORD

@@ -3,9 +3,9 @@ nshutils/actsave/__init__.py,sha256=6gJ49011Ad3kS8BejeZRPj7ZyVyIcmX-VKLEmYZyGM8,
 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=tL-6XvdvJEEr7bje9DSmUBpm6pnJS9XG_1fuw3U1eME,2573
-nshutils/snoop.py,sha256=XhcFic6jbRP1EaBfcpcdtPdh-4kG6uHWKyZIRN4EPnM,7391
+nshutils/logging.py,sha256=B-S4fPKbMPjJwMrEBrA83VvsBca_3sAokSmzn0A0EEw,3384
+nshutils/snoop.py,sha256=JW2JZg6oZd9WWubOP6ok6GUbh5mIQUOdP_T29iH5RCE,7369
 nshutils/typecheck.py,sha256=_KtfinRy9A0Dgq78kN5MGGrefvb6jn2tGY6svdLEBAs,4915
-nshutils-0.12.2.dist-info/METADATA,sha256=w7k32VmAqHvikP9CCkaWartLJ2EOe3UpqC8U20AN8M0,4073
-nshutils-0.12.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshutils-0.12.2.dist-info/RECORD,,
+nshutils-0.14.0.dist-info/METADATA,sha256=59Cz9XX3JUg22c_arXyxQRPl7wHIaGHpbxNxbcjxdZs,4117
+nshutils-0.14.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshutils-0.14.0.dist-info/RECORD,,