Repository: yt-project/unyt
Pull request #512 (2024-07-03 09:26:28)

BUG: fix incompatibilities with NumPy 2.1
- **BUG: test/implement array functions new in NumPy 2.1 (np.cumulative_sum and np.cumulative_prod)**
- **BUG: adapt np.clip's implementation signature at runtime**

[example logs](https://github.com/yt-project/unyt/actions/runs/9774860392/job/26984040518)
xref: https://github.com/numpy/numpy/pull/26724


Patch:
diff --git a/docs/usage.rst b/docs/usage.rst
index a1a7d37..d12be77 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -269,10 +269,9 @@ controlled identically to NumPy arrays, using ``numpy.setprintoptions``:
   >>> import numpy as np
   >>> import unyt as u
   ...
-  >>> np.set_printoptions(precision=4)
-  >>> print([1.123456789]*u.km)
+  >>> with np.printoptions(precision=4):
+  ...     print([1.123456789]*u.km)
   [1.1235] km
-  >>> np.set_printoptions(precision=8)
 
 Print a :math:`\rm{\LaTeX}` representation of a set of units using the
 :meth:`unyt.unit_object.Unit.latex_representation` function or
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index bd752d7..7d6bcb3 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -696,6 +696,13 @@ def cumprod(a, *args, **kwargs):
     )
 
 
+if NUMPY_VERSION >= Version("2.1.0.dev0"):
+
+    @implements(np.cumulative_prod)
+    def cumulative_prod(x, /, *args, **kwargs):
+        return cumprod(x, *args, **kwargs)
+
+
 @implements(np.pad)
 def pad(array, *args, **kwargs):
     return np.pad._implementation(np.asarray(array), *args, **kwargs) * array.units
@@ -816,8 +823,7 @@ def sinc(x, *args, **kwargs):
     return np.sinc._implementation(np.asarray(x), *args, **kwargs)
 
 
-@implements(np.clip)
-def clip(a, a_min, a_max, out=None, *args, **kwargs):
+def clip_impl(a, a_min, a_max, out=None, *args, **kwargs):
     _validate_units_consistency_v2(a.units, a_min, a_max)
     if out is None:
         return (
@@ -843,6 +849,20 @@ def clip(a, a_min, a_max, out=None, *args, **kwargs):
     return unyt_array(res, a.units, bypass_validation=True)
 
 
+if NUMPY_VERSION >= Version("2.1.0.dev0"):
+    from numpy._globals import _NoValue
+
+    @implements(np.clip)
+    def clip(a, a_min=_NoValue, a_max=_NoValue, out=None, *args, **kwargs):
+        return clip_impl(a, a_min, a_max, out, *args, **kwargs)
+
+else:
+
+    @implements(np.clip)
+    def clip(a, a_min, a_max, out=None, *args, **kwargs):
+        return clip_impl(a, a_min, a_max, out, *args, **kwargs)
+
+
 @implements(np.where)
 def where(condition, *args, **kwargs):
     if len(args) == 0:


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index a5960a9..50a59a7 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -187,6 +187,7 @@ if NUMPY_VERSION >= Version("2.0.0dev0"):
 if NUMPY_VERSION >= Version("2.1.0dev0"):
     NOOP_FUNCTIONS |= {
         np.unstack,
+        np.cumulative_sum,
     }
 
 # Functions for which behaviour is intentionally left to default
@@ -1528,28 +1529,44 @@ def test_deltas(func, input_units, output_units):
 
 
 @pytest.mark.parametrize(
-    "func",
+    "func_name",
     [
-        np.cumsum,
-        np.nancumsum,
+        "cumsum",
+        "nancumsum",
+        pytest.param(
+            "cumulative_sum",
+            marks=pytest.mark.skipif(
+                NUMPY_VERSION < Version("2.1.0dev0"),
+                reason="np.cumulative_sum is new in NumPy 2.1",
+            ),
+        ),
     ],
 )
-def test_cumsum(func):
+def test_cumsum(func_name):
     a = [1, 2, 3] * cm
+    func = getattr(np, func_name)
     res = func(a)
     assert type(res) is unyt_array
     assert res.units == cm
 
 
 @pytest.mark.parametrize(
-    "func",
+    "func_name",
     [
-        np.cumprod,
-        np.nancumprod,
+        "cumprod",
+        "nancumprod",
+        pytest.param(
+            "cumulative_prod",
+            marks=pytest.mark.skipif(
+                NUMPY_VERSION < Version("2.1.0dev0"),
+                reason="np.cumulative_prod is new in NumPy 2.1",
+            ),
+        ),
     ],
 )
-def test_cumprod(func):
+def test_cumprod(func_name):
     a = [1, 2, 3] * cm
+    func = getattr(np, func_name)
     with pytest.raises(
         UnytError,
         match=re.escape(