nshutils 0.15.0__tar.gz → 0.16.0__tar.gz

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

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

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

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

{nshutils-0.15.0 → nshutils-0.16.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.16.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)