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

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.2.dist-info → nshutils-0.22.3.dist-info}/METADATA

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

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

@@ -11,12 +11,12 @@ nshutils/lovely/_base.py,sha256=c2XxNJlEdET2mP2gLzlYY1KHsEN4H9eDD_x8SptuBTA,4277
 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.2.dist-info/METADATA,sha256=Fyf30EVxNxg10Vrrg3lEXnTXA5aByv2UBtMqBS6LEg4,4406
-nshutils-0.22.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-nshutils-0.22.2.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,,