Repository: yt-project/unyt
Pull request #500 (2024-04-02 18:13:02)

BUG: fix signature incompatibilities in NEP 18 wrapped functions
I ported a meta test that I originally wrote for astropy.units, and fixed a bunch of subtle signature incompatibilites it detected.
xref https://github.com/astropy/astropy/pull/15710


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 5eebb4d..83fb406 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -21,7 +21,7 @@ repos:
     - id: check-toml
 
 -   repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.1.1
     hooks:
     - id: black
 
@@ -32,7 +32,7 @@ repos:
         additional_dependencies: [black==24.1.1]
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.3.5
+  rev: v0.2.0
   hooks:
   - id: ruff
     args: [--fix, --show-fixes]
diff --git a/docs/extensions/show_all_units.py b/docs/extensions/show_all_units.py
index e9507dc..b4cbee0 100644
--- a/docs/extensions/show_all_units.py
+++ b/docs/extensions/show_all_units.py
@@ -175,14 +175,22 @@ def as_rest_table(data, full=False):
     start_of_line = start_of_line.translate(to_separator)
     vertical_separator = vertical_separator.translate(to_separator)
     end_of_line = end_of_line.translate(to_separator)
-    separator = f"{start_of_line}{vertical_separator.join([x * line_marker for x in sizes])}{end_of_line}"
+    separator = "{}{}{}".format(
+        start_of_line,
+        vertical_separator.join([x * line_marker for x in sizes]),
+        end_of_line,
+    )
     # determine header separator
     th_separator_tr = maketrans("-", "=")
     start_of_line = start_of_line.translate(th_separator_tr)
     line_marker = line_marker.translate(th_separator_tr)
     vertical_separator = vertical_separator.translate(th_separator_tr)
     end_of_line = end_of_line.translate(th_separator_tr)
-    th_separator = f"{start_of_line}{vertical_separator.join([x * line_marker for x in sizes])}{end_of_line}"
+    th_separator = "{}{}{}".format(
+        start_of_line,
+        vertical_separator.join([x * line_marker for x in sizes]),
+        end_of_line,
+    )
     # prepare result
     table.append(separator)
     # set table header
diff --git a/unyt/_array_functions.py b/unyt/_array_functions.py
index 4834fbc..b8322d6 100644
--- a/unyt/_array_functions.py
+++ b/unyt/_array_functions.py
@@ -272,57 +272,57 @@ def cross(a, b, *args, **kwargs):
 
 
 @implements(np.intersect1d)
-def intersect1d(arr1, arr2, /, assume_unique=False, return_indices=False):
-    _validate_units_consistency((arr1, arr2))
+def intersect1d(ar1, ar2, assume_unique=False, return_indices=False):
+    _validate_units_consistency((ar1, ar2))
     retv = np.intersect1d._implementation(
-        np.asarray(arr1),
-        np.asarray(arr2),
+        np.asarray(ar1),
+        np.asarray(ar2),
         assume_unique=assume_unique,
         return_indices=return_indices,
     )
     if return_indices:
         return retv
     else:
-        return retv * arr1.units
+        return retv * ar1.units
 
 
 @implements(np.union1d)
-def union1d(arr1, arr2, /):
-    _validate_units_consistency((arr1, arr2))
-    return np.union1d._implementation(np.asarray(arr1), np.asarray(arr2)) * arr1.units
+def union1d(ar1, ar2):
+    _validate_units_consistency((ar1, ar2))
+    return np.union1d._implementation(np.asarray(ar1), np.asarray(ar2)) * ar1.units
 
 
 @implements(np.linalg.norm)
-def norm(x, /, *args, **kwargs):
+def norm(x, *args, **kwargs):
     return np.linalg.norm._implementation(np.asarray(x), *args, **kwargs) * x.units
 
 
 @implements(np.vstack)
-def vstack(tup, /):
+def vstack(tup, **kwargs):
     ret_units = _validate_units_consistency(tup)
-    return np.vstack._implementation([np.asarray(_) for _ in tup]) * ret_units
+    return np.vstack._implementation([np.asarray(_) for _ in tup], **kwargs) * ret_units
 
 
 @implements(np.hstack)
-def hstack(tup, /):
+def hstack(tup, **kwargs):
     ret_units = _validate_units_consistency(tup)
-    return np.vstack._implementation([np.asarray(_) for _ in tup]) * ret_units
+    return np.vstack._implementation([np.asarray(_) for _ in tup], **kwargs) * ret_units
 
 
 @implements(np.dstack)
-def dstack(tup, /):
+def dstack(tup):
     ret_units = _validate_units_consistency(tup)
     return np.dstack._implementation([np.asarray(_) for _ in tup]) * ret_units
 
 
 @implements(np.column_stack)
-def column_stack(tup, /):
+def column_stack(tup):
     ret_units = _validate_units_consistency(tup)
     return np.column_stack._implementation([np.asarray(_) for _ in tup]) * ret_units
 
 
 @implements(np.stack)
-def stack(arrays, /, axis=0, out=None):
+def stack(arrays, axis=0, out=None, **kwargs):
     ret_units = _validate_units_consistency(arrays)
     if out is None:
         return (
@@ -330,7 +330,7 @@ def stack(arrays, /, axis=0, out=None):
             * ret_units
         )
     res = np.stack._implementation(
-        [np.asarray(_) for _ in arrays], axis=axis, out=np.asarray(out)
+        [np.asarray(_) for _ in arrays], axis=axis, out=np.asarray(out), **kwargs
     )
     if getattr(out, "units", None) is not None:
         out.units = ret_units
@@ -678,8 +678,8 @@ def diff(a, *args, **kwargs):
 
 
 @implements(np.ediff1d)
-def ediff1d(a, *args, **kwargs):
-    return diff_helper(np.ediff1d, a, *args, **kwargs)
+def ediff1d(ary, *args, **kwargs):
+    return diff_helper(np.ediff1d, ary, *args, **kwargs)
 
 
 @implements(np.ptp)
@@ -874,7 +874,8 @@ def tril(m, *args, **kwargs):
 
 
 @implements(np.einsum)
-def einsum(subscripts, *operands, out=None, **kwargs):
+def einsum(*operands, out=None, **kwargs):
+    subscripts, *operands = operands
     ret_units = _validate_units_consistency(operands)
 
     if out is not None:
diff --git a/unyt/_parsing.py b/unyt/_parsing.py
index f22d2ef..3c5f06d 100644
--- a/unyt/_parsing.py
+++ b/unyt/_parsing.py
@@ -85,6 +85,9 @@ def parse_unyt_expr(unit_expr):
             unit_expr, global_dict=global_dict, transformations=unit_text_transform
         )
     except Exception as e:
-        msg = f"Unit expression '{unit_expr}' raised an error during parsing:\n{e!r}"
+        msg = "Unit expression '{}' raised an error during parsing:\n{}".format(
+            unit_expr,
+            repr(e),
+        )
         raise UnitParseError(msg)
     return unit_expr


Test patch:
diff --git a/unyt/tests/test_array_functions.py b/unyt/tests/test_array_functions.py
index 077f999..c39b58b 100644
--- a/unyt/tests/test_array_functions.py
+++ b/unyt/tests/test_array_functions.py
@@ -1,4 +1,5 @@
 # tests for NumPy __array_function__ support
+import inspect
 import re
 from importlib.metadata import version
 
@@ -24,6 +25,13 @@ from unyt.unit_registry import UnitRegistry
 
 NUMPY_VERSION = Version(version("numpy"))
 
+
+VAR_POSITIONAL = inspect.Parameter.VAR_POSITIONAL
+VAR_KEYWORD = inspect.Parameter.VAR_KEYWORD
+POSITIONAL_ONLY = inspect.Parameter.POSITIONAL_ONLY
+KEYWORD_ONLY = inspect.Parameter.KEYWORD_ONLY
+POSITIONAL_OR_KEYWORD = inspect.Parameter.POSITIONAL_OR_KEYWORD
+
 # this is a subset of NOT_HANDLED_FUNCTIONS for which there's nothing to do
 # because they don't apply to (real) numeric types
 # or they work as expected out of the box
@@ -1965,3 +1973,146 @@ def test_interp():
     res = np.interp(_x * cm, _xp * cm, _fp * K)
     assert type(res) is unyt_array
     assert res.units == K
+
+
+@pytest.mark.parametrize(
+    "target, helper",
+    sorted(
+        HANDLED_FUNCTIONS.items(),
+        key=lambda items: items[0].__name__,
+    ),
+    ids=lambda func: func.__name__,
+)
+class TestFunctionHelpersSignatureCompatibility:
+    """
+    Check that a helper function's signature is *at least* as flexible
+    as the helped (target) function's. E.g., any argument that is allowed positionally,
+    or as keyword, by the target must be re-exposed *somehow* by the helper.
+    We explicitly allow helper's signature to be *more* flexible than the target signature
+    by allowing *args and **kwargs catch-all arguments, which we use to limit code
+    duplication, and also help with forward and backward compatibility.
+    See https://github.com/astropy/astropy/issues/15703
+    """
+
+    # this test class is adapted from astropy.units
+
+    @staticmethod
+    def have_catchall_argument(parameters, kind) -> bool:
+        return any(p.kind is kind for p in parameters.values())
+
+    @staticmethod
+    def get_param_group(parameters, kinds: list) -> list[str]:
+        return [name for name, p in parameters.items() if p.kind in kinds]
+
+    def test_all_arguments_reexposed(self, target, helper):
+        try:
+            sig_target = inspect.signature(target)
+        except ValueError:
+            pytest.skip("Non Python function cannot be inspected at runtime")
+
+        params_target = sig_target.parameters
+        sig_helper = inspect.signature(helper)
+        params_helper = sig_helper.parameters
+
+        have_args_helper = self.have_catchall_argument(params_helper, VAR_POSITIONAL)
+        have_kwargs_helper = self.have_catchall_argument(params_helper, VAR_KEYWORD)
+
+        args_helper = list(params_helper.items())
+
+        pos_helper = 0
+        for nt, pt in params_target.items():
+            kt = pt.kind
+            if kt in (POSITIONAL_ONLY, POSITIONAL_OR_KEYWORD):
+                assert pos_helper < len(args_helper), (
+                    "helper's signature is too short; "
+                    "some arguments are not properly re-exposed"
+                )
+                nh, ph = args_helper[pos_helper]
+                if (kh := ph.kind) is not VAR_POSITIONAL:
+                    assert nh == nt, f"argument {nt!r} isn't re-exposed as positional"
+                    assert kh is kt, (
+                        f"helper is not re-exposing argument {nt!r} properly:"
+                        f"expected {kt}, got {kh}"
+                    )
+                    pos_helper += 1
+                    continue
+
+            if kt in (KEYWORD_ONLY, POSITIONAL_OR_KEYWORD):
+                if nt in params_helper:
+                    kh = params_helper[nt].kind
+                    assert kh is kt, (
+                        f"helper is not re-exposing argument {nt!r} properly: "
+                        f"expected {kt}, got {kh}"
+                    )
+                elif kt is KEYWORD_ONLY:
+                    assert (
+                        have_kwargs_helper
+                    ), f"argument {nt!r} is not re-exposed as keyword"
+                elif kt is POSITIONAL_OR_KEYWORD:
+                    assert (
+                        have_args_helper and have_kwargs_helper
+                    ), f"argument {nt!r} is not re-exposed as positional-or-keyword"
+            elif kt is VAR_POSITIONAL:
+                assert have_args_helper, "helper is missing a catch-all *args argument"
+            elif kt is VAR_KEYWORD:
+                assert (
+                    have_kwargs_helper
+                ), "helper is missing a catch-all **kwargs argument"
+
+    def test_known_arguments(self, target, helper):
+        # validate that all exposed arguments map to something in the target
+        try:
+            sig_target = inspect.signature(target)
+        except ValueError:
+            pytest.skip("Non Python function cannot be inspected at runtime")
+
+        params_target = sig_target.parameters
+        sig_helper = inspect.signature(helper)
+        params_helper = sig_helper.parameters
+
+        for kind in (POSITIONAL_ONLY, POSITIONAL_OR_KEYWORD):
+            args_target = self.get_param_group(params_helper, [kind])
+            args_helper = self.get_param_group(params_helper, [kind])
+
+            if (nhelper := len(args_helper)) > (ntarget := len(args_target)):
+                unknown: list[str] = args_helper[ntarget:]
+                raise AssertionError(
+                    f"Found unknown {kind} parameter(s) "
+                    "in helper's signature: "
+                    f"{unknown}, at position(s) {list(range(ntarget, nhelper))}"
+                )
+
+        # keyword-allowed
+        keyword_allowed_target = set(
+            self.get_param_group(params_target, [KEYWORD_ONLY, POSITIONAL_OR_KEYWORD])
+        )
+        keyword_allowed_helper = set(
+            self.get_param_group(params_helper, [KEYWORD_ONLY, POSITIONAL_OR_KEYWORD])
+        )
+
+        # additional private keyword-only argument are allowed because
+        # they are only intended for testing purposes.
+        # For instance, quantile has such a parameter '_q_unit'
+        keyword_allowed_helper = {
+            name for name in keyword_allowed_helper if not name.startswith("_")
+        }
+
+        diff = keyword_allowed_helper - keyword_allowed_target
+        assert not diff, (
+            "Found some keyword-allowed parameters in helper "
+            f"that are unknown to target: {diff}"
+        )
+
+        # finally, check that default values are correctly replicated
+        for name, ph in params_helper.items():
+            if name not in params_target:
+                # In a few cases, the helper defines names that are not in
+                # the target (e.g., a private name like _q_unit in quantile,
+                # or a *args, **kwargs that captures further arguments
+                # that do not matter. We let such cases slip by.
+                continue
+            pt = params_target[name]
+            assert ph.default == pt.default, (
+                f"Default value mismatch for argument {name!r}. "
+                f"Helper has {ph.default!r}, target has {pt.default!r}"
+            )
diff --git a/unyt/tests/test_unyt_array.py b/unyt/tests/test_unyt_array.py
index cdce53c..52b881b 100644
--- a/unyt/tests/test_unyt_array.py
+++ b/unyt/tests/test_unyt_array.py
@@ -1546,8 +1546,7 @@ def test_h5_io():
 
     reg.add("code_length", 10.0, dimensions.length)
 
-    rng = np.random.default_rng()
-    warr = unyt_array(rng.random((256, 256)), "code_length", registry=reg)
+    warr = unyt_array(np.random.random((256, 256)), "code_length", registry=reg)
 
     warr.write_hdf5("test.h5")
 
@@ -1566,8 +1565,8 @@ def test_h5_io():
 
     # test code to overwrite existing dataset with data that has a different
     # shape
-    rng = np.random.default_rng()
-    warr = unyt_array(rng.random((255, 255)), "code_length", registry=reg)
+
+    warr = unyt_array(np.random.random((255, 255)), "code_length", registry=reg)
 
     warr.write_hdf5("test.h5")
 
@@ -2166,10 +2165,9 @@ def test_loadtxt_and_savetxt():
     curdir = os.getcwd()
     os.chdir(tmpdir)
 
-    rng = np.random.default_rng()
-    a = unyt_array(rng.random(10), "kpc")
-    b = unyt_array(rng.random(10), "Msun")
-    c = unyt_array(rng.random(10), "km/s")
+    a = unyt_array(np.random.random(10), "kpc")
+    b = unyt_array(np.random.random(10), "Msun")
+    c = unyt_array(np.random.random(10), "km/s")
 
     savetxt("arrays.dat", [a, b, c], delimiter=",")
 
@@ -2223,9 +2221,8 @@ def test_loadtxt_and_savetxt():
 
 
 def test_trig_ufunc_degrees():
-    rng = np.random.default_rng()
     for ufunc in (np.sin, np.cos, np.tan):
-        degree_values = rng.random(10) * degree
+        degree_values = np.random.random(10) * degree
         radian_values = degree_values.in_units("radian")
         assert_array_equal(ufunc(degree_values), ufunc(radian_values))