nshutils 0.15.0__tar.gz → 0.17.0__tar.gz

{nshutils-0.15.0 → nshutils-0.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshutils
-Version: 0.15.0
+Version: 0.17.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshutils-0.15.0 → nshutils-0.17.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshutils"
-version = "0.15.0"
+version = "0.17.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshutils-0.15.0 → nshutils-0.17.0}/src/nshutils/__init__.py

@@ -1,8 +1,12 @@
 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 .logging import init_python_logging as init_python_logging
 from .logging import lovely as lovely
 from .logging import pretty as pretty
 from .snoop import snoop as snoop
+from .typecheck import tassert as tassert
+from .typecheck import typecheck_modules as typecheck_modules
+from .typecheck import typecheck_this_module as typecheck_this_module

nshutils-0.17.0/src/nshutils/_display.py

@@ -0,0 +1,60 @@
+import importlib.util
+from functools import cache
+from typing import Any
+
+
+@cache
+def _in_ipython():
+    try:
+        from IPython import get_ipython
+
+        return get_ipython() is not None
+    except ImportError:
+        return False
+
+
+@cache
+def _treescope_installed():
+    return importlib.util.find_spec("treescope") is not None
+
+
+@cache
+def _rich_installed():
+    return importlib.util.find_spec("rich") is not None
+
+
+def display(*args: Any):
+    """
+        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
+    installed, it will fall back to IPython's display function. In a
+    non-IPython environment, rich will be used if available, otherwise
+    the standard print function will be used.
+
+    Args:
+        *args: Any objects to display.
+    """
+    if _in_ipython():
+        if _treescope_installed():
+            import treescope
+
+            with treescope.active_autovisualizer.set_scoped(
+                treescope.ArrayAutovisualizer()
+            ):
+                treescope.display(*args)
+        elif _rich_installed():
+            import rich
+
+            rich.print(*args)
+        else:
+            from IPython.display import display
+
+            display(*args)
+    elif _rich_installed():
+        import rich
+
+        rich.print(*args)
+    else:
+        print(*args)

{nshutils-0.15.0 → nshutils-0.17.0}/src/nshutils/logging.py

@@ -19,7 +19,7 @@ def init_python_logging(
 
             _lovely_tensors.monkey_patch()
         except ImportError:
-            logging.exception(
+            logging.info(
                 "Failed to import `lovely_tensors`. Ignoring pretty PyTorch tensor formatting"
             )
 
@@ -29,7 +29,7 @@ def init_python_logging(
 
             _lovely_numpy.set_config(repr=_lovely_numpy.lovely)
         except ImportError:
-            logging.exception(
+            logging.info(
                 "Failed to import `lovely_numpy`. Ignoring pretty numpy array formatting"
             )
 
@@ -45,11 +45,11 @@ def init_python_logging(
                     autovisualize_arrays=treescope_autovisualize_arrays
                 )
             else:
-                logging.exception(
+                logging.info(
                     "Treescope setup is only supported in Jupyter notebooks. Skipping."
                 )
         except ImportError:
-            logging.exception(
+            logging.info(
                 "Failed to import `treescope` or `IPython`. Ignoring `treescope` registration"
             )
 
@@ -65,7 +65,7 @@ def init_python_logging(
 
             log_handlers.append(RichHandler(rich_tracebacks=rich_tracebacks))
         except ImportError:
-            logging.exception(
+            logging.info(
                 "Failed to import rich. Falling back to default Python logging."
             )