nshutils 0.13.0__py3-none-any.whl → 0.15.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,
@@ -17,7 +19,7 @@ def init_python_logging(
 
             _lovely_tensors.monkey_patch()
         except ImportError:
-            logging.warning(
+            logging.exception(
                 "Failed to import `lovely_tensors`. Ignoring pretty PyTorch tensor formatting"
             )
 
@@ -27,10 +29,30 @@ def init_python_logging(
 
             _lovely_numpy.set_config(repr=_lovely_numpy.lovely)
         except ImportError:
-            logging.warning(
+            logging.exception(
                 "Failed to import `lovely_numpy`. Ignoring pretty numpy array formatting"
             )
 
+    if treescope:
+        try:
+            # Check if we're in a Jupyter environment
+            from IPython import get_ipython
+
+            if get_ipython() is not None:
+                import treescope as _treescope  # type: ignore
+
+                _treescope.basic_interactive_setup(
+                    autovisualize_arrays=treescope_autovisualize_arrays
+                )
+            else:
+                logging.exception(
+                    "Treescope setup is only supported in Jupyter notebooks. Skipping."
+                )
+        except ImportError:
+            logging.exception(
+                "Failed to import `treescope` or `IPython`. Ignoring `treescope` registration"
+            )
+
     log_handlers: list[logging.Handler] = []
     if log_save_dir:
         log_file = log_save_dir / "logging.log"
@@ -43,7 +65,7 @@ def init_python_logging(
 
             log_handlers.append(RichHandler(rich_tracebacks=rich_tracebacks))
         except ImportError:
-            logging.warning(
+            logging.exception(
                 "Failed to import rich. Falling back to default Python logging."
             )
 
@@ -59,14 +81,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 +104,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-0.13.0.dist-info → nshutils-0.15.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.13.0
+Version: 0.15.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.13.0.dist-info → nshutils-0.15.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/logging.py,sha256=3NwW2pSoXBktiRWPGtUzs8BrOW8H9ZpfI-8wiLbEnEE,3721
 nshutils/snoop.py,sha256=JW2JZg6oZd9WWubOP6ok6GUbh5mIQUOdP_T29iH5RCE,7369
 nshutils/typecheck.py,sha256=_KtfinRy9A0Dgq78kN5MGGrefvb6jn2tGY6svdLEBAs,4915
-nshutils-0.13.0.dist-info/METADATA,sha256=gby5UkK98NiS-G0xv5jhIh8CY5pHYLH2cd6UHWzepbU,4073
-nshutils-0.13.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshutils-0.13.0.dist-info/RECORD,,
+nshutils-0.15.0.dist-info/METADATA,sha256=0Ak2KBfkCHz2g-8ytTX2YIm6VF5Jic_HJT3JG0t8F0U,4117
+nshutils-0.15.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshutils-0.15.0.dist-info/RECORD,,