Repository: xgcm/xgcm
Pull request #441 (2022-02-23 16:54:48)

Signature class + refactor inside apply_as_grid_ufunc
Refactors grid_ufunc internals by splitting things up into smaller functions. Also adds a `Signature` class, which is neat, but I'm not sure it needs to be exposed to users.

 - [x] Passes `pre-commit run --all-files`


Patch:
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 1214dc8..f498c02 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -14,8 +14,8 @@ from . import comodo, gridops
 from .duck_array_ops import _apply_boundary_condition, _pad_array, concatenate
 from .grid_ufunc import (
     GridUFunc,
+    Signature,
     _has_chunked_core_dims,
-    _signatures_equivalent,
     apply_as_grid_ufunc,
 )
 from .metrics import iterate_axis_combinations
@@ -27,7 +27,13 @@ try:
 except ImportError:
     numba = None
 
-from typing import Any, Dict, Iterable, Union
+from typing import Any, Callable, Dict, Iterable, List, Mapping, Tuple, Union
+
+# Only need this until python 3.8
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal  # type: ignore
 
 docstrings = docrep.DocstringProcessor(doc_key="My doc string")
 
@@ -1811,7 +1817,7 @@ class Grid:
 
         return self._transpose_to_keep_same_dim_order(da, array, axis)
 
-    def _create_1d_grid_ufunc_signatures(self, da, axis, to):
+    def _create_1d_grid_ufunc_signatures(self, da, axis, to) -> List[Signature]:
         """
         Create a list of signatures to pass to apply_grid_ufunc.
 
@@ -1828,7 +1834,7 @@ class Grid:
             if to_pos is None:
                 to_pos = ax._default_shifts[from_pos]
 
-            signature_1d = f"({ax_name}:{from_pos})->({ax_name}:{to_pos})"
+            signature_1d = Signature(f"({ax_name}:{from_pos})->({ax_name}:{to_pos})")
             signatures.append(signature_1d)
 
         return signatures
@@ -1854,13 +1860,14 @@ class Grid:
 
     def apply_as_grid_ufunc(
         self,
-        func,
-        *args,
-        signature="",
-        boundary_width=None,
-        boundary=None,
-        fill_value=None,
-        dask="forbidden",
+        func: Callable,
+        *args: xr.DataArray,
+        signature: Union[str, Signature] = "",
+        boundary_width: Mapping[str, Tuple[int, int]] = None,
+        boundary: Union[str, Mapping[str, str]] = None,
+        fill_value: Union[float, Mapping[str, float]] = None,
+        dask: Literal["forbidden", "parallelized", "allowed"] = "forbidden",
+        map_overlap: bool = False,
         **kwargs,
     ):
         """
@@ -1926,6 +1933,7 @@ class Grid:
             boundary=boundary,
             fill_value=fill_value,
             dask=dask,
+            map_overlap=map_overlap,
             **kwargs,
         )
 
@@ -2348,7 +2356,7 @@ class Grid:
         return self._apply_vector_function(Axis.diff, vector, **kwargs)
 
 
-def _select_grid_ufunc(funcname, signature, module, **kwargs):
+def _select_grid_ufunc(funcname, signature: Signature, module, **kwargs):
     # TODO to select via other kwargs (e.g. boundary) the signature of this function needs to be generalised
 
     def is_grid_ufunc(obj):
@@ -2366,9 +2374,7 @@ def _select_grid_ufunc(funcname, signature, module, **kwargs):
         )
 
     signature_matching_ufuncs = [
-        f
-        for f in name_matching_ufuncs
-        if _signatures_equivalent(f.signature, signature)
+        f for f in name_matching_ufuncs if f.signature.equivalent(signature)
     ]
     if len(signature_matching_ufuncs) == 0:
         raise NotImplementedError(
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index 5529551..9497c92 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -7,11 +7,18 @@ from typing import (
     Dict,
     List,
     Mapping,
+    Optional,
     Sequence,
     Tuple,
     Union,
 )
 
+# Only need this until python 3.8
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal  # type: ignore
+
 import numpy as np
 import xarray as xr
 
@@ -97,6 +104,66 @@ def _parse_grid_ufunc_signature(
     return in_ax_names, out_ax_names, in_ax_pos, out_ax_pos
 
 
+class Signature:
+    signature: str
+    in_ax_names: List[Tuple[str, ...]]
+    out_ax_names: List[Tuple[str, ...]]
+    in_ax_positions: List[Tuple[str, ...]]
+    out_ax_positions: List[Tuple[str, ...]]
+
+    _REPLACEMENT_DUMMY_INDEX_NAMES = [f"__{char}" for char in string.ascii_letters]
+
+    def __init__(self, signature: str):
+        self.signature: str = signature
+        (
+            self.in_ax_names,
+            self.out_ax_names,
+            self.in_ax_positions,
+            self.out_ax_positions,
+        ) = _parse_grid_ufunc_signature(signature)
+
+    def __str__(self):
+        return self.signature
+
+    def __repr__(self):
+        return f"Signature({self.signature})"
+
+    def equivalent(self, other: "Signature") -> bool:
+        """Whether or not two signatures are equivalent."""
+
+        """
+        Axes names in signatures are dummy variables, so an exact string match is not required.
+        Our comparison strategy is to instead work through both signatures left to right, replacing all occurrences
+        of each dummy index with names drawn from a common list. If after this process the replaced names are not
+        identical, the signatures must not be equivalent. Axes positions do have to match exactly.
+        """
+
+        def set_unique_inds(sig_part):
+            return set([i for arg in sig_part for i in arg])
+
+        all_unique_sig1_indices = set_unique_inds(self.in_ax_names) | set_unique_inds(
+            self.out_ax_names
+        )
+        all_unique_sig2_indices = set_unique_inds(other.in_ax_names) | set_unique_inds(
+            other.out_ax_names
+        )
+
+        if len(all_unique_sig1_indices) != len(all_unique_sig2_indices):
+            return False
+
+        sig1_replaced = self.signature
+        sig2_replaced = other.signature
+        for dummy1, dummy2, common_replacement in zip(
+            all_unique_sig1_indices,
+            all_unique_sig2_indices,
+            self._REPLACEMENT_DUMMY_INDEX_NAMES,
+        ):
+            sig1_replaced = sig1_replaced.replace(dummy1, common_replacement)
+            sig2_replaced = sig2_replaced.replace(dummy2, common_replacement)
+
+        return sig1_replaced == sig2_replaced
+
+
 class GridUFunc:
     """
     Binds a numpy ufunc into a "grid-aware ufunc", meaning that when called ufunc is wrapped by `apply_as_grid_ufunc`.
@@ -142,9 +209,16 @@ class GridUFunc:
     Grid.apply_as_grid_ufunc
     """
 
+    ufunc: Callable
+    signature: Signature
+    boundary_width: Optional[Mapping[str, Tuple[int, int]]]
+    dask: Literal["forbidden", "parallelized", "allowed"]
+    map_overlap: bool
+
     def __init__(self, ufunc: Callable, **kwargs):
-        self.ufunc = ufunc
-        self.signature = kwargs.pop("signature", "")
+        self.ufunc = ufunc  # type: ignore  # see mypy issue 2427
+        signature = kwargs.pop("signature", "")
+        self.signature = Signature(signature)
         self.boundary_width = kwargs.pop("boundary_width", None)
         self.dask = kwargs.pop("dask", "forbidden")
         self.map_overlap = kwargs.pop("map_overlap", False)
@@ -246,12 +320,12 @@ def apply_as_grid_ufunc(
     *args: xr.DataArray,
     axis: Sequence[str],
     grid: "Grid" = None,
-    signature: str = "",
+    signature: Union[str, Signature] = "",
     boundary_width: Mapping[str, Tuple[int, int]] = None,
     boundary: Union[str, Mapping[str, str]] = None,
     fill_value: Union[float, Mapping[str, float]] = None,
     keep_coords: bool = True,
-    dask: str = "forbidden",
+    dask: Literal["forbidden", "parallelized", "allowed"] = "forbidden",
     map_overlap: bool = False,
     **kwargs,
 ) -> List[Any]:
@@ -345,25 +419,23 @@ def apply_as_grid_ufunc(
         )
 
     # Extract Axes information from signature
-    (
-        in_dummy_ax_names,
-        out_dummy_ax_names,
-        in_ax_pos,
-        out_ax_pos,
-    ) = _parse_grid_ufunc_signature(signature)
+    if not isinstance(signature, Signature):
+        sig = Signature(signature)
+    else:
+        sig = signature
 
     dummy_to_real_axes_mapping = _identify_dummy_axes_with_real_axes(
-        in_dummy_ax_names, axis
+        sig.in_ax_names, axis
     )
 
     # Determine names of output axes from names in signature
     # TODO what if we need to add a new core dim to the output that does match an input axis? Where do we get the name from?
     out_ax_names = [
-        [dummy_to_real_axes_mapping[ax] for ax in arg] for arg in out_dummy_ax_names
+        [dummy_to_real_axes_mapping[ax] for ax in arg] for arg in sig.out_ax_names
     ]
 
     # Check that input args are in correct grid positions
-    for i, (arg_ns, arg_ps, arg) in enumerate(zip(axis, in_ax_pos, args)):
+    for i, (arg_ns, arg_ps, arg) in enumerate(zip(axis, sig.in_ax_positions, args)):
         for n, p in zip(arg_ns, arg_ps):
             try:
                 ax_pos = grid.axes[n].coords[p]
@@ -384,11 +456,11 @@ def apply_as_grid_ufunc(
     # Determine core dimensions for apply_ufunc
     in_core_dims = [
         [grid.axes[n].coords[p] for n, p in zip(arg_ns, arg_ps)]
-        for arg_ns, arg_ps in zip(axis, in_ax_pos)
+        for arg_ns, arg_ps in zip(axis, sig.in_ax_positions)
     ]
     out_core_dims = [
         [grid.axes[n].coords[p] for n, p in zip(arg_ns, arg_ps)]
-        for arg_ns, arg_ps in zip(out_ax_names, out_ax_pos)
+        for arg_ns, arg_ps in zip(out_ax_names, sig.out_ax_positions)
     ]
 
     # TODO allow users to specify new output dtypes
@@ -434,56 +506,17 @@ def apply_as_grid_ufunc(
         rechunked_padded_args = padded_args
 
     if map_overlap:
-        # map operation over dask chunks along core dimensions
-        from dask.array import map_overlap as dask_map_overlap  # type: ignore
-
-        # Need to transpose the numpy axis arguments to leave core dims at end
-        # else they won't match up inside mapped_func after xr.apply_ufunc does its transposition
-        transposed_original_args = [
-            arg.transpose(..., *in_core_dims[i]) for i, arg in enumerate(args)
-        ]
-
-        boundary_width_per_numpy_axis = {
-            grid.axes[ax_name]._get_axis_dim_num(transposed_original_args[0]): width
-            for ax_name, width in boundary_width_real_axes.items()
-        }
-
         # Disallow situations where shifting axis position would cause chunk size to change
-        _check_if_length_would_change(out_dummy_ax_names, in_ax_pos, out_ax_pos)
-
-        single_dim_chunktype = Tuple[int, ...]
-
-        def _dict_to_numbered_axes(
-            sizes: Mapping[str, single_dim_chunktype]
-        ) -> Tuple[single_dim_chunktype, ...]:
-            """This implicitly crystallises the order of the given mapping"""
-            return tuple(sizes.values())
-
-        # Our rechunking means dask.map_overlap needs to be explicitly told what chunks output should have
-        # But in this case output chunks are the same as input chunks
-        # (as we disallowed axis positions for which this is not the case)
-        original_chunksizes = [
-            arg.variable.chunksizes for arg in transposed_original_args
-        ]
-        # TODO first argument only because map_overlap can't handle multiple return values (I think)
-        true_chunksizes = original_chunksizes[0]
-        # dask.map_overlap needs chunks in terms of axis number, not axis name (i.e. (chunks, ...), not {str: chunks})
-        true_chunksizes_per_numpy_axis = _dict_to_numbered_axes(true_chunksizes)
-
-        # (we don't need a separate code path using bare map_blocks if boundary_widths are zero because map_overlap just
-        # calls map_blocks automatically in that scenario)
-        def mapped_func(*a, **kw):
-            return dask_map_overlap(
-                func,
-                *a,
-                **kw,
-                depth=boundary_width_per_numpy_axis,
-                boundary="none",
-                trim=False,
-                meta=np.array([], dtype=out_dtypes[0]),
-                chunks=true_chunksizes_per_numpy_axis,
-            )
+        _check_if_length_would_change(sig)
 
+        mapped_func = _map_func_over_core_dims(
+            func,
+            args,
+            grid,
+            in_core_dims,
+            boundary_width_real_axes,
+            out_dtypes,
+        )
     else:
         mapped_func = func
 
@@ -513,36 +546,9 @@ def apply_as_grid_ufunc(
         results = (results,)
 
     # Restore any dimension coordinates associated with new output dims that are present in grid
-    results_with_coords = []
-    for res, arg_out_core_dims in zip(results, out_core_dims):
-
-        # Only reconstruct coordinates that actually contain grid position info (i.e. not just integer values along a dim.)
-        # Therefore if input only had dimensions and no coordinates, the output should too.
-        new_core_dim_coords = {
-            dim: grid._ds.coords[dim]
-            for dim in arg_out_core_dims
-            if dim in grid._ds.coords and dim not in res.coords
-        }
-
-        try:
-            res = res.assign_coords(new_core_dim_coords)
-        except ValueError as err:
-            if boundary_width and str(err).startswith("conflicting sizes"):
-                # TODO make this error more informative?
-                raise ValueError(
-                    f"{str(err)} - does your grid ufunc correctly trim off the same number of elements "
-                    f"which were added by padding using boundary_width={boundary_width}?"
-                )
-            else:
-                raise
-
-        if not keep_coords:
-            # TODO I don't like the `keep_coords` argument in general and think it should be removed for clarity.
-            # Drop any non-dimension coordinates on the output
-            non_dim_coords = [coord for coord in res.coords if coord not in res.dims]
-            res = res.drop_vars(non_dim_coords)
-
-        results_with_coords.append(res)
+    results_with_coords = _reattach_coords(
+        results, grid, out_core_dims, boundary_width, keep_coords
+    )
 
     # Return single results not wrapped in 1-element tuple, like xr.apply_ufunc does
     if len(results_with_coords) == 1:
@@ -553,34 +559,93 @@ def apply_as_grid_ufunc(
     return results_with_coords
 
 
-def is_dim_chunked(a, dim):
+def _is_dim_chunked(a, dim):
     # TODO this func can't handle Datasets - it will error if you check multiple variables with different chunking
     return len(a.variable.chunksizes[dim]) > 1
 
 
 def _has_chunked_core_dims(obj: xr.DataArray, core_dims: Sequence[str]) -> bool:
     # TODO what if only some of the core dimensions are chunked?
-    return obj.chunks is not None and any(is_dim_chunked(obj, dim) for dim in core_dims)
+    return obj.chunks is not None and any(
+        _is_dim_chunked(obj, dim) for dim in core_dims
+    )
+
+
+def _map_func_over_core_dims(
+    func,
+    args,
+    grid,
+    in_core_dims,
+    boundary_width_real_axes,
+    out_dtypes,
+):
+    """Map operation over dask chunks along core dimensions"""
+
+    from dask.array import map_overlap as dask_map_overlap  # type: ignore
+
+    # Need to transpose the numpy axis arguments to leave core dims at end
+    # else they won't match up inside mapped_func after xr.apply_ufunc does its transposition
+    transposed_original_args = [
+        arg.transpose(..., *in_core_dims[i]) for i, arg in enumerate(args)
+    ]
+
+    boundary_width_per_numpy_axis = {
+        grid.axes[ax_name]._get_axis_dim_num(transposed_original_args[0]): width
+        for ax_name, width in boundary_width_real_axes.items()
+    }
+
+    single_dim_chunktype = Tuple[int, ...]
+
+    def _dict_to_numbered_axes(
+        sizes: Mapping[str, single_dim_chunktype]
+    ) -> Tuple[single_dim_chunktype, ...]:
+        """This implicitly crystallises the order of the given mapping"""
+        return tuple(sizes.values())
+
+    # Our rechunking means dask.map_overlap needs to be explicitly told what chunks output should have
+    # But in this case output chunks are the same as input chunks
+    # (as we disallowed axis positions for which this is not the case)
+    original_chunksizes = [arg.variable.chunksizes for arg in transposed_original_args]
+    # TODO first argument only because map_overlap can't handle multiple return values (I think)
+    true_chunksizes = original_chunksizes[0]
+    # dask.map_overlap needs chunks in terms of axis number, not axis name (i.e. (chunks, ...), not {str: chunks})
+    true_chunksizes_per_numpy_axis = _dict_to_numbered_axes(true_chunksizes)
+
+    # (we don't need a separate code path using bare map_blocks if boundary_widths are zero because map_overlap just
+    # calls map_blocks automatically in that scenario)
+    def mapped_func(*a, **kw):
+        return dask_map_overlap(
+            func,
+            *a,
+            **kw,
+            depth=boundary_width_per_numpy_axis,
+            boundary="none",
+            trim=False,
+            meta=np.array([], dtype=out_dtypes[0]),
+            chunks=true_chunksizes_per_numpy_axis,
+        )
+
+    return mapped_func
 
 
 DISALLOWED_OVERLAP_POSITIONS = ["inner", "outer"]
 
 
-def _check_if_length_would_change(
-    out_ax_names: List[Tuple[str, ...]],
-    in_ax_pos: List[Tuple[str, ...]],
-    out_ax_pos: List[Tuple[str, ...]],
-):
+def _check_if_length_would_change(signature: Signature):
     """Check if map_overlap can actually handle the complexity of this signature."""
 
     # TODO this restriction is because dask.array.map_overlap does not currently allow for multiple return arrays
-    if len(out_ax_names) > 1:
+    if len(signature.out_ax_names) > 1:
         raise NotImplementedError(
             "Currently cannot automatically map a ufunc over multiple outputs when the core "
             "dimension is chunked"
         )
 
-    all_ax_positions = set(p for arg_ps in in_ax_pos + out_ax_pos for p in arg_ps)
+    all_ax_positions = set(
+        p
+        for arg_ps in signature.in_ax_positions + signature.out_ax_positions
+        for p in arg_ps
+    )
     if any(pos in DISALLOWED_OVERLAP_POSITIONS for pos in all_ax_positions):
         raise NotImplementedError(
             "Cannot chunk along a core dimension for a grid ufunc which has a signature which "
@@ -688,35 +753,36 @@ def _identify_dummy_axes_with_real_axes(
     return dict(zip(unique_dummy_axes, unique_real_axes))
 
 
-_REPLACEMENT_DUMMY_INDEX_NAMES = [f"__{char}" for char in string.ascii_letters]
-
-
-def _signatures_equivalent(sig1: str, sig2: str) -> bool:
-    """
-    Axes names in signatures are dummy variables, so an exact string match is not required.
+def _reattach_coords(results, grid, out_core_dims, boundary_width, keep_coords):
+    results_with_coords = []
+    for res, arg_out_core_dims in zip(results, out_core_dims):
 
-    Our comparison strategy is to instead work through both signatures left to right, replacing all occurrences
-    of each dummy index with names drawn from a common list. If after this process the replaced names are not
-    identical, the signatures must not be equivalent. Axes positions do have to match exactly.
-    """
-    sig1_in, sig1_out, _, _ = _parse_grid_ufunc_signature(sig1)
-    sig2_in, sig2_out, _, _ = _parse_grid_ufunc_signature(sig2)
+        # Only reconstruct coordinates that actually contain grid position info (i.e. not just integer values along a dim.)
+        # Therefore if input only had dimensions and no coordinates, the output should too.
+        new_core_dim_coords = {
+            dim: grid._ds.coords[dim]
+            for dim in arg_out_core_dims
+            if dim in grid._ds.coords and dim not in res.coords
+        }
 
-    all_unique_sig1_indices = set([i for arg in sig1_in for i in arg])
-    all_unique_sig2_indices = set([i for arg in sig2_in for i in arg])
+        try:
+            res = res.assign_coords(new_core_dim_coords)
+        except ValueError as err:
+            if boundary_width and str(err).startswith("conflicting sizes"):
+                # TODO make this error more informative?
+                raise ValueError(
+                    f"{str(err)} - does your grid ufunc correctly trim off the same number of elements "
+                    f"which were added by padding using boundary_width={boundary_width}?"
+                )
+            else:
+                raise
 
-    if len(all_unique_sig1_indices) != len(all_unique_sig2_indices):
-        return False
+        if not keep_coords:
+            # TODO I don't like the `keep_coords` argument in general and think it should be removed for clarity.
+            # Drop any non-dimension coordinates on the output
+            non_dim_coords = [coord for coord in res.coords if coord not in res.dims]
+            res = res.drop_vars(non_dim_coords)
 
-    sig1_replaced = sig1
-    sig2_replaced = sig2
-    for dummy1, dummy2, common_replacement in zip(
-        all_unique_sig1_indices, all_unique_sig2_indices, _REPLACEMENT_DUMMY_INDEX_NAMES
-    ):
-        sig1_replaced = sig1_replaced.replace(dummy1, common_replacement)
-        sig2_replaced = sig2_replaced.replace(dummy2, common_replacement)
+        results_with_coords.append(res)
 
-    if sig1_replaced == sig2_replaced:
-        return True
-    else:
-        return False
+    return results_with_coords


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index d480744..bad9a48 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -9,8 +9,8 @@ from xarray.testing import assert_equal
 from xgcm.grid import Grid, _select_grid_ufunc
 from xgcm.grid_ufunc import (
     GridUFunc,
+    Signature,
     _parse_grid_ufunc_signature,
-    _signatures_equivalent,
     apply_as_grid_ufunc,
     as_grid_ufunc,
 )
@@ -171,7 +171,7 @@ class TestGridUFunc:
             return a - np.roll(a, shift=-1)
 
         assert isinstance(diff_center_to_left, GridUFunc)
-        assert diff_center_to_left.signature == signature
+        assert str(diff_center_to_left.signature) == signature
 
         with pytest.raises(TypeError, match="Unsupported keyword argument"):
 
@@ -741,27 +741,27 @@ class TestMapOverlapGridops:
 
 class TestSignaturesEquivalent:
     def test_equivalent(self):
-        sig1 = "(X:center)->(X:left)"
-        sig2 = "(X:center)->(X:left)"
-        assert _signatures_equivalent(sig1, sig2)
+        sig1 = Signature("(X:center)->(X:left)")
+        sig2 = Signature("(X:center)->(X:left)")
+        assert sig1.equivalent(sig2)
 
-        sig3 = "(Y:center)->(Y:left)"
-        assert _signatures_equivalent(sig1, sig3)
+        sig3 = Signature("(Y:center)->(Y:left)")
+        assert sig1.equivalent(sig3)
 
     def test_not_equivalent(self):
-        sig1 = "(X:center)->(X:left)"
-        sig2 = "(X:center)->(X:center)"
-        assert not _signatures_equivalent(sig1, sig2)
+        sig1 = Signature("(X:center)->(X:left)")
+        sig2 = Signature("(X:center)->(X:center)")
+        assert not sig1.equivalent(sig2)
 
-        sig3 = "(X:center)->(Y:left)"
-        assert not _signatures_equivalent(sig1, sig3)
+        sig3 = Signature("(X:center)->(Y:left)")
+        assert not sig1.equivalent(sig3)
 
-        sig4 = "(X:center,X:center)->(X:left)"
-        assert not _signatures_equivalent(sig1, sig4)
+        sig4 = Signature("(X:center,X:center)->(X:left)")
+        assert not sig1.equivalent(sig4)
 
     def test_no_indices(self):
-        sig = "()->()"
-        assert _signatures_equivalent(sig, sig)
+        sig = Signature("()->()")
+        assert sig.equivalent(sig)
 
 
 class GridOpsMockUp:
@@ -794,34 +794,41 @@ class GridOpsMockUp:
 class TestGridUFuncDispatch:
     def test_select_ufunc(self):
         gridufunc, _ = _select_grid_ufunc(
-            "diff", "(X:center)->(X:left)", module=GridOpsMockUp
+            "diff", Signature("(X:center)->(X:left)"), module=GridOpsMockUp
         )
         assert gridufunc is GridOpsMockUp.diff_center_to_left
 
     def test_select_ufunc_equivalent_signature(self):
         gridufunc, _ = _select_grid_ufunc(
-            "diff", "(Y:center)->(Y:left)", module=GridOpsMockUp
+            "diff", Signature("(Y:center)->(Y:left)"), module=GridOpsMockUp
         )
         assert gridufunc is GridOpsMockUp.diff_center_to_left
 
         with pytest.raises(NotImplementedError):
-            _select_grid_ufunc("diff", "(X:center)->(Y:left)", module=GridOpsMockUp)
+            _select_grid_ufunc(
+                "diff", Signature("(X:center)->(Y:left)"), module=GridOpsMockUp
+            )
 
     def test_select_ufunc_wrong_signature(self):
         with pytest.raises(NotImplementedError):
-            _select_grid_ufunc("diff", "(X:center)->(X:center)", module=GridOpsMockUp)
+            _select_grid_ufunc(
+                "diff", Signature("(X:center)->(X:center)"), module=GridOpsMockUp
+            )
 
     @pytest.mark.xfail(reason="currently no need for this")
     def test_select_ufunc_by_kwarg(self):
         gridufunc, _ = _select_grid_ufunc(
-            "diff", "(X:center)->(X:right)", module=GridOpsMockUp, boundary="fill"
+            "diff",
+            Signature("(X:center)->(X:right)"),
+            module=GridOpsMockUp,
+            boundary="fill",
         )
         assert gridufunc is GridOpsMockUp.diff_center_to_right_fill
 
         with pytest.raises(NotImplementedError):
             _select_grid_ufunc(
                 "diff",
-                "(X:center)->(X:right)",
+                Signature("(X:center)->(X:right)"),
                 module=GridOpsMockUp,
                 boundary="nonsense",
             )
@@ -830,6 +837,6 @@ class TestGridUFuncDispatch:
     def test_pass_through_other_kwargs(self):
         # TODO put this in test_grid.py instead?
         gridufunc, _ = _select_grid_ufunc(
-            "pass", "()->()", module=GridOpsMockUp, boundary="fill"
+            "pass", Signature("()->()"), module=GridOpsMockUp, boundary="fill"
         )
         assert gridufunc(a=1) == {"a": 1}