nshutils 0.3.0__tar.gz → 0.4.0__tar.gz

{nshutils-0.3.0 → nshutils-0.4.0}/PKG-INFO

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

{nshutils-0.3.0 → nshutils-0.4.0}/pyproject.toml

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

nshutils-0.4.0/src/nshutils/__init__.py

@@ -0,0 +1,6 @@
+from . import actsave as actsave
+from . import typecheck as typecheck
+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

nshutils-0.4.0/src/nshutils/logging.py

@@ -0,0 +1,88 @@
+import logging
+from pathlib import Path
+
+
+def init_python_logging(
+    *,
+    lovely_tensors: bool = False,
+    lovely_numpy: bool = False,
+    rich: bool = False,
+    log_level: int | str | None = logging.INFO,
+    log_save_dir: Path | None = None,
+):
+    if lovely_tensors:
+        try:
+            import lovely_tensors as _lovely_tensors
+
+            _lovely_tensors.monkey_patch()
+        except ImportError:
+            logging.warning(
+                "Failed to import `lovely_tensors`. Ignoring pretty PyTorch tensor formatting"
+            )
+
+    if lovely_numpy:
+        try:
+            import lovely_numpy as _lovely_numpy
+
+            _lovely_numpy.set_config(repr=_lovely_numpy.lovely)
+        except ImportError:
+            logging.warning(
+                "Failed to import `lovely_numpy`. Ignoring pretty numpy array formatting"
+            )
+
+    log_handlers: list[logging.Handler] = []
+    if log_save_dir:
+        log_file = log_save_dir / "logging.log"
+        log_file.touch(exist_ok=True)
+        log_handlers.append(logging.FileHandler(log_file))
+
+    if rich:
+        try:
+            from rich.logging import RichHandler
+
+            log_handlers.append(RichHandler())
+        except ImportError:
+            logging.warning(
+                "Failed to import rich. Falling back to default Python logging."
+            )
+
+    logging.basicConfig(
+        level=log_level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=log_handlers,
+    )
+
+
+def pretty(
+    *,
+    lovely_tensors: bool = True,
+    lovely_numpy: bool = True,
+    log_level: int | str | None = logging.INFO,
+    log_save_dir: Path | None = None,
+    rich_log_handler: bool = True,
+):
+    init_python_logging(
+        lovely_tensors=lovely_tensors,
+        lovely_numpy=lovely_numpy,
+        rich=rich_log_handler,
+        log_level=log_level,
+        log_save_dir=log_save_dir,
+    )
+
+
+def lovely(
+    *,
+    lovely_tensors: bool = True,
+    lovely_numpy: bool = True,
+    log_level: int | str | None = logging.INFO,
+    log_save_dir: Path | None = None,
+    rich_log_handler: bool = True,
+):
+    pretty(
+        lovely_tensors=lovely_tensors,
+        lovely_numpy=lovely_numpy,
+        log_level=log_level,
+        log_save_dir=log_save_dir,
+        rich_log_handler=rich_log_handler,
+    )

nshutils-0.3.0/src/nshutils/__init__.py

@@ -1,3 +0,0 @@
-from . import actsave as actsave
-from . import typecheck as typecheck
-from .snoop import snoop as snoop