nshutils 0.22.1__py3-none-any.whl → 0.22.3__py3-none-any.whl

nshutils/lovely/_monkey_patch_all.py

@@ -41,21 +41,20 @@ def monkey_patch(libraries: list[Library] | Literal["auto"] = "auto"):
 
     with contextlib.ExitStack() as stack:
         for library in libraries:
-            match library:
-                case "torch":
-                    from .torch_ import torch_monkey_patch
+            if library == "torch":
+                from .torch_ import torch_monkey_patch
 
-                    stack.enter_context(torch_monkey_patch())
-                case "jax":
-                    from .jax_ import jax_monkey_patch
+                stack.enter_context(torch_monkey_patch())
+            elif library == "jax":
+                from .jax_ import jax_monkey_patch
 
-                    stack.enter_context(jax_monkey_patch())
-                case "numpy":
-                    from .numpy_ import numpy_monkey_patch
+                stack.enter_context(jax_monkey_patch())
+            elif library == "numpy":
+                from .numpy_ import numpy_monkey_patch
 
-                    stack.enter_context(numpy_monkey_patch())
-                case _:
-                    assert_never(library)
+                stack.enter_context(numpy_monkey_patch())
+            else:
+                assert_never(library)
 
         log.info(
             f"Monkey patched libraries: {', '.join(libraries)}. "

nshutils/lovely/numpy_.py

@@ -8,6 +8,18 @@ from ._base import lovely_repr, monkey_patch_contextmanager
 from .utils import LovelyStats, array_stats
 
 
+def _np_ge_2():
+    import importlib.metadata
+
+    from packaging.version import Version
+
+    try:
+        numpy_version = importlib.metadata.version("numpy")
+        return Version(numpy_version) >= Version("2.0")
+    except importlib.metadata.PackageNotFoundError:
+        return False
+
+
 def _type_name(array: np.ndarray):
     return (
         "array"
@@ -55,18 +67,43 @@ def numpy_repr(array: np.ndarray) -> LovelyStats:
     }
 
 
-@monkey_patch_contextmanager(dependencies=["numpy"])
-def numpy_monkey_patch():
-    try:
-        np.set_printoptions(override_repr=numpy_repr)
-        logging.info(
-            f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
-            f"{np.get_printoptions()=}"
-        )
-        yield
-    finally:
-        np.set_printoptions(override_repr=None)
-        logging.info(
-            f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
-            f"{np.get_printoptions()=}"
-        )
+# If numpy 2.0, use the new API override_repr.
+if _np_ge_2():
+
+    @monkey_patch_contextmanager(dependencies=["numpy"])
+    def numpy_monkey_patch():
+        try:
+            np.set_printoptions(override_repr=numpy_repr)
+            logging.info(
+                f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"{np.get_printoptions()=}"
+            )
+            yield
+        finally:
+            np.set_printoptions(override_repr=None)
+            logging.info(
+                f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"{np.get_printoptions()=}"
+            )
+
+else:
+
+    @monkey_patch_contextmanager(dependencies=["numpy"])
+    def numpy_monkey_patch():
+        try:
+            np.set_string_function(numpy_repr, True)  # pyright: ignore[reportAttributeAccessIssue]
+            np.set_string_function(numpy_repr, False)  # pyright: ignore[reportAttributeAccessIssue]
+
+            logging.info(
+                f"Numpy monkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"{np.get_printoptions()=}"
+            )
+            yield
+        finally:
+            np.set_string_function(None, True)  # pyright: ignore[reportAttributeAccessIssue]
+            np.set_string_function(None, False)  # pyright: ignore[reportAttributeAccessIssue]
+
+            logging.info(
+                f"Numpy unmonkey patching: using {numpy_repr.__name__} for numpy arrays. "
+                f"{np.get_printoptions()=}"
+            )

{nshutils-0.22.1.dist-info → nshutils-0.22.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshutils
-Version: 0.22.1
+Version: 0.22.3
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshutils-0.22.1.dist-info → nshutils-0.22.3.dist-info}/RECORD

@@ -8,15 +8,15 @@ nshutils/display.py,sha256=Ge63yllx7gi-MKL3mKQeQ5doql_nj56-o5aoTVmusDg,1473
 nshutils/logging.py,sha256=78pv3-I_gmbKSf5_mYYBr6_H4GNBGErghAdhH9wfYIc,2205
 nshutils/lovely/__init__.py,sha256=gbWMNs7xfK1CiNdkHvfH0KcyaGjdZ8_WUBGfaEUDN4I,451
 nshutils/lovely/_base.py,sha256=c2XxNJlEdET2mP2gLzlYY1KHsEN4H9eDD_x8SptuBTA,4277
-nshutils/lovely/_monkey_patch_all.py,sha256=WZsC6Xp5-Z2GBd6xyZZEsD4C2xNyZy0YBfjZzAy3m8M,2028
+nshutils/lovely/_monkey_patch_all.py,sha256=zgMupp2Wc_O9R3arl-BAIePpvQSi6TCeshGMaui-Cc8,1986
 nshutils/lovely/config.py,sha256=lVNMuU1oUvsYlGN0Sn-m6iOLbJIchVnWDpyHm09nWo8,1224
 nshutils/lovely/jax_.py,sha256=mPH-tSOzWE27ymupllBnO4O6avvJJxpF2a1-G4dN5Ow,2214
-nshutils/lovely/numpy_.py,sha256=iHA4BCJIW9IU6DXKEfYbh9RA2xyeXeL0tyHdEja89Sw,1866
+nshutils/lovely/numpy_.py,sha256=CI6pidaohK0tFw9mmSBTaDfjCKrFhY87ga0RIqTqSFk,3162
 nshutils/lovely/torch_.py,sha256=3wnXLa-1xwuQVk1fM50mBqOVDv5wZHjfnBwzOnjcFjg,2638
 nshutils/lovely/utils.py,sha256=2ksT5YGVViFuWc8jSkwVCsABripJmyVJdEDDH7aab70,10459
 nshutils/snoop.py,sha256=7d7_Q5sJmINL1J29wcnxEvpV95zvZYNoVn5frCq-rww,7393
 nshutils/typecheck.py,sha256=Gi7xtfilN_UwZ1FTFqBVKDhcQzBEDonVxIv3bUj-uXY,5582
 nshutils/util.py,sha256=tx-XiRbOrpafV3OkJDE5IVFtzn3kN7uSZ8FkMor0H5c,2845
-nshutils-0.22.1.dist-info/METADATA,sha256=GHCw6WkX4yi-KGaDXf-EVgee6AfYnUPD4Z452Wq14rc,4406
-nshutils-0.22.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-nshutils-0.22.1.dist-info/RECORD,,
+nshutils-0.22.3.dist-info/METADATA,sha256=Or1chA_NrLS05hiiqbXaj_ph8tZ91TwRT1A5iBfTOW8,4406
+nshutils-0.22.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+nshutils-0.22.3.dist-info/RECORD,,