Repository: xgcm/xgcm
Pull request #418 (2022-02-03 18:34:56)

Grid information through type annotations
 (Replacement for #417 without the massive diff)

 - [x] Implements the [suggestion for specifying variable grid information in type hints](https://github.com/xgcm/xgcm/issues/344#issuecomment-1020334503)
 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`

We have to upgrade to python 3.9 to use `typing.Annotated`.

TomNicholas: I have a weird mypy/pre-commit problem, where if I do `pre-commit run --all-files` then mypy passes, but once I actually try to commit anything I get mypy errors:
```
Trim Trailing Whitespace.................................................Passed
Fix End of Files.........................................................Passed
Check Yaml...........................................(no files to check)Skipped
isort....................................................................Passed
black....................................................................Passed
blackdoc.................................................................Passed
flake8...................................................................Passed
mypy.....................................................................Failed
- hook id: mypy
- exit code: 1

xgcm/grid_ufunc.py:6: error: Module "typing" has no attribute "Annotated"
xgcm/grid_ufunc.py:202: error: Unexpected keyword argument "include_extras" for "get_type_hints"
Found 2 errors in 1 file (checked 1 source file)
```
To me those errors look like something to do with the version of python or `typing_extensions`, because `Annotated` and `include_extras` were only added in python 3.9. However (a) I thought I updated all that and (b) I don't understand how pre-commit can lead to different results in these two cases. For the moment I'm just forcing all the commits through with `--no-verify`.
TomNicholas: > A bit torn about making 3.9 minimum

Xarray is already on 3.8 minimum, and they are fairly conservative, so we're hardly pushing all our users into a brave new world!
jbusecke: > Xarray is already on 3.8 minimum, and they are fairly conservative, so we're hardly pushing all our users into a brave new world!

Sure, then l have nothing really against it. @rabernat ?
TomNicholas: > backwards compatibility

Also in general I personally don't really think of version changes as "breaking backwards compatibility". At the most updating our minimum version requires the user only to update their conda environment, whereas when I think of "breaking backwards compatibility" I imagine changes that actually require users to update their actual production code. Obviously no changes at all would be best, but one is less of a big deal than the other imo.
jbusecke: Fair point.

Patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index d533f76..196025f 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -23,7 +23,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.7", "3.8", "3.9"]
+        python-version: ["3.9", "3.10"]
     steps:
       - uses: actions/checkout@v2
       - name: Cache conda
@@ -79,7 +79,7 @@ jobs:
       with:
         activate-environment: test_env_xgcm # Defined in ci/environment-core-deps.yml
         auto-update-conda: false
-        python-version: 3.8
+        python-version: 3.9
         environment-file: ci/environment-core-deps.yml
         use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
     - name: Set up conda environment
@@ -112,7 +112,7 @@ jobs:
       with:
         activate-environment: test_env_xgcm # Defined in ci/environment-upstream-dev.yml
         auto-update-conda: false
-        python-version: 3.8
+        python-version: 3.9
         environment-file: ci/environment-upstream-dev.yml
         use-only-tar-bz2: true # IMPORTANT: This needs to be set for caching to work properly!
     - name: Set up conda environment
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a5ae150..00f095c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -34,7 +34,7 @@ repos:
             # Type stubs
             types-pkg_resources,
             types-PyYAML,
-            typing-extensions==3.10.0.0,
+            typing-extensions==4.0.1,
             # Dependencies that are typed
             numpy,
           ]
diff --git a/ci/environment-core-deps.yml b/ci/environment-core-deps.yml
index ea419ba..71c832d 100644
--- a/ci/environment-core-deps.yml
+++ b/ci/environment-core-deps.yml
@@ -14,6 +14,6 @@ dependencies:
   - pytest
   - pytest-cov
   - pytest-xdist
-  - python=3.8
+  - python=3.9
   - xarray
   - typing_extensions
diff --git a/doc/environment.yml b/doc/environment.yml
index 0a3b747..8f19bd6 100644
--- a/doc/environment.yml
+++ b/doc/environment.yml
@@ -1,6 +1,6 @@
 name: test_env_xgcm_docs
 dependencies:
-  - python=3.6
+  - python=3.9
   - xarray
   - netcdf4
   - cartopy
diff --git a/setup.cfg b/setup.cfg
index d7a5bae..6be15cd 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -15,6 +15,8 @@ ignore =
     W503
     E265
     F811
+    # Allows type hinting as Gridded[DataArray, "(X:center)"], where we did `from typing import Annotated as Gridded`
+    F722
 
 [isort]
 profile = black
diff --git a/setup.py b/setup.py
index 2121986..790f3b1 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@ install_requires = [
     "numpy",
     "future",
     "docrep<=0.2.7",
-    "typing_extensions >= 3.7",  # python_version < '3.8'
+    "typing_extensions >= 4.0.1",  # python_version < '3.8'
 ]
 doc_requires = [
     "sphinx",
@@ -48,14 +48,12 @@ setup(
         "Topic :: Scientific/Engineering",
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
-        "Programming Language :: Python :: 3.7",
-        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
     ],
     packages=find_packages(exclude=["docs", "tests", "tests.*", "docs.*"]),
     install_requires=install_requires,
     extras_require=extras_require,
-    python_requires=">=3.7",
+    python_requires=">=3.9",
     long_description=long_description,
     long_description_content_type="text/x-rst",
     setup_requires="setuptools_scm",
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 5d5f9b0..2353c96 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -14,7 +14,7 @@ from . import comodo, gridops
 from .duck_array_ops import _apply_boundary_condition, _pad_array, concatenate
 from .grid_ufunc import (
     GridUFunc,
-    Signature,
+    _GridUFuncSignature,
     _has_chunked_core_dims,
     apply_as_grid_ufunc,
 )
@@ -1778,7 +1778,8 @@ class Grid:
         else:
             dask = "forbidden"
 
-        array = da
+        # Need to copy to avoid modifying in-place. Ideally we would test for this behaviour specifically
+        array = da.copy(deep=False)
         # Apply 1D function over multiple axes
         # TODO This will call xarray.apply_ufunc once for each axis, but if signatures + kwargs are the same then we
         # TODO only actually need to call apply_ufunc once for those axes
@@ -1817,7 +1818,9 @@ class Grid:
 
         return self._transpose_to_keep_same_dim_order(da, array, axis)
 
-    def _create_1d_grid_ufunc_signatures(self, da, axis, to) -> List[Signature]:
+    def _create_1d_grid_ufunc_signatures(
+        self, da, axis, to
+    ) -> List[_GridUFuncSignature]:
         """
         Create a list of signatures to pass to apply_grid_ufunc.
 
@@ -1834,7 +1837,10 @@ class Grid:
             if to_pos is None:
                 to_pos = ax._default_shifts[from_pos]
 
-            signature_1d = Signature(f"({ax_name}:{from_pos})->({ax_name}:{to_pos})")
+            # TODO build this more directly?
+            signature_1d = _GridUFuncSignature.from_string(
+                f"({ax_name}:{from_pos})->({ax_name}:{to_pos})"
+            )
             signatures.append(signature_1d)
 
         return signatures
@@ -1863,7 +1869,7 @@ class Grid:
         func: Callable,
         *args: xr.DataArray,
         axis: Sequence[Sequence[str]] = None,
-        signature: Union[str, Signature] = "",
+        signature: Union[str, _GridUFuncSignature] = "",
         boundary_width: Mapping[str, Tuple[int, int]] = None,
         boundary: Union[str, Mapping[str, str]] = None,
         fill_value: Union[float, Mapping[str, float]] = None,
@@ -2366,7 +2372,7 @@ class Grid:
         return self._apply_vector_function(Axis.diff, vector, **kwargs)
 
 
-def _select_grid_ufunc(funcname, signature: Signature, module, **kwargs):
+def _select_grid_ufunc(funcname, signature: _GridUFuncSignature, 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):
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index 9604adc..105e41a 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -6,19 +6,15 @@ from typing import (
     Callable,
     Dict,
     List,
+    Literal,
     Mapping,
     Optional,
     Sequence,
     Tuple,
     Union,
+    get_type_hints,
 )
 
-# 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
 
@@ -31,6 +27,7 @@ if TYPE_CHECKING:
 # Modified version of `numpy.lib.function_base._parse_gufunc_signature`
 # Modifications:
 #   - Specify xgcm.Axis name and "axis positions" instead of numpy axes as (ax_name:ax_pos)
+
 _AXIS_NAME = r"\w+"
 _AXIS_POSITION = "(?:center|left|right|inner|outer)"
 _AXIS_NAME_POSITION_PAIR = f"{_AXIS_NAME}:{_AXIS_POSITION}"
@@ -42,96 +39,96 @@ _ARGUMENT_LIST = f"{_ARGUMENT}(?:,{_ARGUMENT})*"
 _SIGNATURE = f"^{_ARGUMENT_LIST}->{_ARGUMENT_LIST}$"
 
 
-def _parse_grid_ufunc_signature(
-    signature: str,
-) -> Tuple[
-    List[Tuple[str, ...]],
-    List[Tuple[str, ...]],
-    List[Tuple[str, ...]],
-    List[Tuple[str, ...]],
-]:
-    """
-    Parse string signatures for a grid-aware universal function.
-
-    The way this parser works excludes using Axis names that match possible xgcm
-    axis positions, i.e. ['center', 'left', 'right', 'inner', 'outer'].
+T_AX_POS_LIST = List[Tuple[str, ...]]
 
-    Arguments
-    ---------
-    signature : string
-        Grid universal function signature. Specifies the xgcm.Axis names and
-        positions for each input and output variable, e.g.,
 
-        ``"(X:center)->(X:left)"`` for ``diff_center_to_left(a)`.
+class _GridUFuncSignature:
+    """
+    Core xGCM Axes and grid positions signature for a given function.
 
-    Returns
-    -------
-    input_axes_names : List[Tuple[str, ...]]
-        Input Axes names parsed from the signature
-    output_axes_names : List[Tuple[str, ...]]
-        Output Axes names parsed from the signature
-    input_axes_positions : List[Tuple[str, ...]]
-        Input Axes positions parsed from the signature
-    output_axes_positions : List[Tuple[str, ...]]
-        Output Axes positions parsed from the signature
+    Based on the signature provided by generalized ufuncs in NumPy, and in xarray.
     """
 
-    signature = signature.replace(" ", "")
+    in_ax_names: T_AX_POS_LIST
+    in_ax_positions: T_AX_POS_LIST
+    out_ax_names: T_AX_POS_LIST
+    out_ax_positions: T_AX_POS_LIST
 
-    if not re.match(_SIGNATURE, signature):
-        raise ValueError(f"Not a valid grid ufunc signature: {signature}")
+    _REPLACEMENT_DUMMY_INDEX_NAMES = [f"__{char}" for char in string.ascii_letters]
 
-    in_txt, out_txt = signature.split("->")
+    def __init__(
+        self,
+        in_ax_names: T_AX_POS_LIST,
+        in_ax_positions: T_AX_POS_LIST,
+        out_ax_names: T_AX_POS_LIST,
+        out_ax_positions: T_AX_POS_LIST,
+    ):
+        """Construct the grid signature directly from its internal attributes."""
 
-    in_ax_names = []
-    for arg in re.findall(_ARGUMENT, in_txt):
-        # Delete the axis positions so they aren't matched as axis names
-        only_names = re.sub(_AXIS_POSITION, "", arg)
-        in_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+        if not in_ax_names or not in_ax_positions:
+            raise ValueError(
+                "At least one input argument of the Grid UFunc signature must have "
+                "axis names and positions"
+            )
+        else:
+            self.in_ax_names = in_ax_names
+            self.in_ax_positions = in_ax_positions
 
-    out_ax_names = []
-    for arg in re.findall(_ARGUMENT, out_txt):
-        only_names = re.sub(_AXIS_POSITION, "", arg)
-        out_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+        # Can imagine grid ufuncs where outputs have no core dimensions (e.g. result of inner product)
+        self.out_ax_names = out_ax_names
+        self.out_ax_positions = out_ax_positions
 
-    in_ax_pos = [
-        tuple(re.findall(_AXIS_POSITION, arg)) for arg in re.findall(_ARGUMENT, in_txt)
-    ]
-    out_ax_pos = [
-        tuple(re.findall(_AXIS_POSITION, arg)) for arg in re.findall(_ARGUMENT, out_txt)
-    ]
+    def __str__(self):
+        """The string representation of this signature object"""
+
+        in_arg_sigs = [
+            ",".join(f"{ax}:{pos}" for ax, pos in zip(arg_in_names, arg_in_pos))
+            for arg_in_names, arg_in_pos in zip(self.in_ax_names, self.in_ax_positions)
+        ]
+        lhs = ",".join(f"({arg_sig})" for arg_sig in in_arg_sigs)
+
+        out_arg_sigs = [
+            ",".join(f"{ax}:{pos}" for ax, pos in zip(arg_out_names, arg_out_pos))
+            for arg_out_names, arg_out_pos in zip(
+                self.out_ax_names, self.out_ax_positions
+            )
+        ]
+        rhs = ",".join(f"({arg_sig})" for arg_sig in out_arg_sigs)
 
-    return in_ax_names, out_ax_names, in_ax_pos, out_ax_pos
+        return f"{lhs}->{rhs}"
 
+    @classmethod
+    def from_string(cls, signature: str) -> "_GridUFuncSignature":
+        """Constructs the grid signature from its string representation."""
+        (
+            in_ax_names,
+            in_ax_positions,
+            out_ax_names,
+            out_ax_positions,
+        ) = _parse_signature_from_string(signature)
 
-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, ...]]
+        return cls(in_ax_names, in_ax_positions, out_ax_names, out_ax_positions)
 
-    _REPLACEMENT_DUMMY_INDEX_NAMES = [f"__{char}" for char in string.ascii_letters]
+    @classmethod
+    def from_type_hints(cls, hints: Dict[str, Any]) -> "_GridUFuncSignature":
+        """
+        Constructs the grid signature from the type hints of a function, and returns it ready for parsing.
 
-    def __init__(self, signature: str):
-        self.signature: str = signature
+        Type hints must first be obtained using `typing.get_type_hints(ufunc, include_extras=True)`.
+        """
         (
-            self.in_ax_names,
-            self.out_ax_names,
-            self.in_ax_positions,
-            self.out_ax_positions,
-        ) = _parse_grid_ufunc_signature(signature)
+            in_ax_names,
+            in_ax_positions,
+            out_ax_names,
+            out_ax_positions,
+        ) = _parse_signature_from_type_hints(hints)
 
-    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."""
+        return cls(in_ax_names, in_ax_positions, out_ax_names, out_ax_positions)
 
+    def equivalent(self, other: "_GridUFuncSignature") -> 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
@@ -151,8 +148,8 @@ class Signature:
         if len(all_unique_sig1_indices) != len(all_unique_sig2_indices):
             return False
 
-        sig1_replaced = self.signature
-        sig2_replaced = other.signature
+        sig1_replaced = str(self)
+        sig2_replaced = str(other)
         for dummy1, dummy2, common_replacement in zip(
             all_unique_sig1_indices,
             all_unique_sig2_indices,
@@ -164,6 +161,113 @@ class Signature:
         return sig1_replaced == sig2_replaced
 
 
+def _parse_signature_from_string(
+    signature: str,
+) -> Tuple[T_AX_POS_LIST, T_AX_POS_LIST, T_AX_POS_LIST, T_AX_POS_LIST]:
+    """
+    Parse string signatures for a grid-aware universal function.
+
+    The way this parser works excludes using Axis names that match possible xgcm
+    axis positions, i.e. ['center', 'left', 'right', 'inner', 'outer'].
+    """
+
+    signature = signature.replace(" ", "")
+
+    if not re.match(_SIGNATURE, signature):
+        raise ValueError(f"Not a valid grid ufunc signature: {signature}")
+
+    in_txt, out_txt = signature.split("->")
+
+    in_ax_names = []
+    for arg in re.findall(_ARGUMENT, in_txt):
+        # Delete the axis positions so they aren't matched as axis names
+        only_names = re.sub(_AXIS_POSITION, "", arg)
+        in_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+
+    out_ax_names = []
+    for arg in re.findall(_ARGUMENT, out_txt):
+        only_names = re.sub(_AXIS_POSITION, "", arg)
+        out_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+
+    in_ax_pos = [
+        tuple(re.findall(_AXIS_POSITION, arg)) for arg in re.findall(_ARGUMENT, in_txt)
+    ]
+    out_ax_pos = [
+        tuple(re.findall(_AXIS_POSITION, arg)) for arg in re.findall(_ARGUMENT, out_txt)
+    ]
+
+    return in_ax_names, in_ax_pos, out_ax_names, out_ax_pos
+
+
+def _parse_signature_from_type_hints(
+    hints: Dict[str, Any]
+) -> Tuple[T_AX_POS_LIST, T_AX_POS_LIST, T_AX_POS_LIST, T_AX_POS_LIST]:
+    """
+    Parse signatures from type annotations for a grid-aware universal function.
+
+    The way this parser works excludes using Axis names that match possible xgcm
+    axis positions, i.e. ['center', 'left', 'right', 'inner', 'outer'].
+    """
+
+    # First do output args
+    try:
+        return_hint = hints.pop("return")
+    except KeyError:
+        # TODO does this cause a problem if the output has >1 return arguments none of which have grid positions?
+        out_ax_names: T_AX_POS_LIST = [()]
+        out_ax_pos: T_AX_POS_LIST = [()]
+    else:
+        return_hints = _maybe_multiple_return_vals(return_hint)
+
+        return_annotations = [
+            hint.__metadata__[0]
+            for hint in return_hints
+            if hasattr(hint, "__metadata__")
+        ]
+
+        out_ax_names = []
+        for arg in return_annotations:
+            # Delete the axis positions so they aren't matched as axis names
+            only_names = re.sub(_AXIS_POSITION, "", arg)
+            out_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+
+        out_ax_pos = [
+            tuple(re.findall(_AXIS_POSITION, arg)) for arg in return_annotations
+        ]
+
+    # Now do input args
+    arg_annotations = [
+        hint.__metadata__[0] for hint in hints.values() if hasattr(hint, "__metadata__")
+    ]
+
+    # TODO check number of annotations?
+
+    in_ax_names = []
+    for arg in arg_annotations:
+        # Delete the axis positions so they aren't matched as axis names
+        only_names = re.sub(_AXIS_POSITION, "", arg)
+        in_ax_names.append(tuple(re.findall(_AXIS_NAME, only_names)))
+
+    in_ax_pos = [tuple(re.findall(_AXIS_POSITION, arg)) for arg in arg_annotations]
+
+    # Do a sanity check before going any further
+    str_signature = str(
+        _GridUFuncSignature(in_ax_names, in_ax_pos, out_ax_names, out_ax_pos)
+    )
+    if not re.match(_SIGNATURE, str_signature):
+        raise ValueError(f"Not a valid grid ufunc signature: {str_signature}")
+
+    return in_ax_names, in_ax_pos, out_ax_names, out_ax_pos
+
+
+def _maybe_multiple_return_vals(return_hint):
+    """if ufunc returns multiple values (each of which might be annotated) we must extract from Tuple first"""
+    return_hints = (
+        list(return_hint.__args__) if return_hint._name == "Tuple" else [return_hint]
+    )
+    return return_hints
+
+
 class GridUFunc:
     """
     Binds a numpy ufunc into a "grid-aware ufunc", meaning that when called ufunc is wrapped by `apply_as_grid_ufunc`.
@@ -210,21 +314,58 @@ class GridUFunc:
     """
 
     ufunc: Callable
-    signature: Signature
+    signature: _GridUFuncSignature
     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  # type: ignore  # see mypy issue 2427
-        signature = kwargs.pop("signature", "")
-        self.signature = Signature(signature)
+
+        str_sig = kwargs.pop("signature")
+        self.signature = self._get_signature_from_str_or_type_hints(ufunc, str_sig)
         self.boundary_width = kwargs.pop("boundary_width", None)
         self.dask = kwargs.pop("dask", "forbidden")
         self.map_overlap = kwargs.pop("map_overlap", False)
         if kwargs:
             raise TypeError("Unsupported keyword argument(s) provided")
 
+    @staticmethod
+    def _get_signature_from_str_or_type_hints(
+        ufunc, str_sig: Optional[str]
+    ) -> _GridUFuncSignature:
+        """Get grid ufunc signature, either from type hints or from string signature kwarg"""
+
+        hints = get_type_hints(ufunc, include_extras=True)
+
+        def _has_annotations(hints):
+            try:
+                # TODO I want this to be .pop but then I get problems with variable scope
+                return_hint = hints["return"]
+            except KeyError:
+                pass
+            else:
+                return_hints = _maybe_multiple_return_vals(return_hint)
+                if any(hasattr(hint, "__metadata__") for hint in return_hints):
+                    return True
+
+            return any(hasattr(hint, "__metadata__") for hint in hints.values())
+
+        if str_sig:
+            if _has_annotations(hints):
+                raise ValueError(
+                    "Must specify axis positions through only one of either type hints or signature kwarg, not both."
+                )
+
+            return _GridUFuncSignature.from_string(str_sig)
+        else:
+            if not _has_annotations(hints):
+                raise ValueError(
+                    "Must specify axis positions through either type hints or signature kwarg"
+                )
+
+            return _GridUFuncSignature.from_type_hints(hints)
+
     def __repr__(self):
         return (
             f"GridUFunc(ufunc={self.ufunc}, signature='{self.signature}', boundary_width='{self.boundary_width}', "
@@ -320,7 +461,7 @@ def apply_as_grid_ufunc(
     *args: xr.DataArray,
     axis: Sequence[Sequence[str]] = None,
     grid: "Grid" = None,
-    signature: Union[str, Signature] = "",
+    signature: Union[str, _GridUFuncSignature] = "",
     boundary_width: Mapping[str, Tuple[int, int]] = None,
     boundary: Union[str, Mapping[str, str]] = None,
     fill_value: Union[float, Mapping[str, float]] = None,
@@ -422,8 +563,8 @@ def apply_as_grid_ufunc(
         )
 
     # Extract Axes information from signature
-    if not isinstance(signature, Signature):
-        sig = Signature(signature)
+    if not isinstance(signature, _GridUFuncSignature):
+        sig = _GridUFuncSignature.from_string(signature)
     else:
         sig = signature
 
@@ -639,7 +780,7 @@ def _map_func_over_core_dims(
 DISALLOWED_OVERLAP_POSITIONS = ["inner", "outer"]
 
 
-def _check_if_length_would_change(signature: Signature):
+def _check_if_length_would_change(signature: _GridUFuncSignature):
     """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


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index e96996e..37a7088 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -1,4 +1,5 @@
 import re
+from typing import Annotated, Tuple
 
 import dask.array  # type: ignore
 import numpy as np
@@ -9,16 +10,16 @@ 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,
+    _GridUFuncSignature,
+    _parse_signature_from_string,
     apply_as_grid_ufunc,
     as_grid_ufunc,
 )
 
 
-class TestParseGridUfuncSignature:
+class TestParseSignatureFromString:
     @pytest.mark.parametrize(
-        "signature, exp_in_ax_names, exp_out_ax_names, exp_in_ax_pos, exp_out_ax_pos",
+        "sig_str, exp_in_ax_names, exp_in_ax_pos, exp_out_ax_names, exp_out_ax_pos",
         [
             ("()->()", [()], [()], [()], [()]),
             ("(X:center)->()", [("X",)], [()], [("center",)], [()]),
@@ -51,18 +52,18 @@ class TestParseGridUfuncSignature:
     )
     def test_parse_valid_signatures(
         self,
-        signature,
+        sig_str,
         exp_in_ax_names,
         exp_out_ax_names,
         exp_in_ax_pos,
         exp_out_ax_pos,
     ):
-        in_ax_names, out_ax_names, in_ax_pos, out_ax_pos = _parse_grid_ufunc_signature(
-            signature
+        in_ax_names, out_ax_names, in_ax_pos, out_ax_pos = _parse_signature_from_string(
+            sig_str
         )
         assert in_ax_names == exp_in_ax_names
-        assert out_ax_names == exp_out_ax_names
         assert in_ax_pos == exp_in_ax_pos
+        assert out_ax_names == exp_out_ax_names
         assert out_ax_pos == exp_out_ax_pos
 
     @pytest.mark.parametrize(
@@ -78,7 +79,140 @@ class TestParseGridUfuncSignature:
     )
     def test_invalid_signatures(self, signature):
         with pytest.raises(ValueError):
-            _parse_grid_ufunc_signature(signature)
+            _parse_signature_from_string(signature)
+
+    @pytest.mark.parametrize(
+        "sig_str",
+        [
+            "()->()",
+            "(X:center)->()",
+            "()->(X:left)",
+            "(X:center)->(X:left)",
+            "(X:left)->(Y:center)",
+            "(X:left)->(Y:center)",
+            "(X:left),(X:right)->(Y:center)",
+            "(X:center)->(Y:inner),(Y:outer)",
+            "(X:center,Y:center)->(Z:center)",
+        ],
+    )
+    def test_roundtrip_from_string(self, sig_str):
+        """Checks that the __str__ method of the signature class works"""
+        sig = _GridUFuncSignature.from_string(sig_str)
+        assert str(sig) == sig_str
+
+
+class TestParseSignatureFromTypeHints:
+    def test_no_args_to_annotate(self):
+        with pytest.raises(ValueError, match="Must specify axis positions"):
+
+            @as_grid_ufunc()
+            def ufunc():
+                ...
+
+    # TODO test hints without annotations
+    # TODO test hints with annotations that don't conform to Xgcm
+
+    def test_annotated_args(self):
+        @as_grid_ufunc()
+        def ufunc(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:center"]:
+            ...
+
+        assert str(ufunc.signature) == "(X:center)->(X:center)"
+
+        @as_grid_ufunc()
+        def ufunc(
+            a: Annotated[np.ndarray, "X:center,Y:center"]
+        ) -> Annotated[np.ndarray, "X:center"]:
+            ...
+
+        assert str(ufunc.signature) == "(X:center,Y:center)->(X:center)"
+
+        @as_grid_ufunc()
+        def ufunc(
+            a: Annotated[np.ndarray, "X:left"],
+            b: Annotated[np.ndarray, "Y:right"],
+        ) -> Annotated[np.ndarray, "X:center"]:
+            ...
+
+        assert str(ufunc.signature) == "(X:left),(Y:right)->(X:center)"
+
+        @as_grid_ufunc()
+        def ufunc(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:left,Y:right"]:
+            ...
+
+        assert str(ufunc.signature) == "(X:center)->(X:left,Y:right)"
+
+        @as_grid_ufunc()
+        def ufunc(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Tuple[Annotated[np.ndarray, "X:left"], Annotated[np.ndarray, "Y:right"]]:
+            ...
+
+        assert str(ufunc.signature) == "(X:center)->(X:left),(Y:right)"
+
+    @pytest.mark.xfail(reason="signature regex will assume nonsense==no inputs")
+    def test_invalid_arg_annotation(self):
+        # TODO how to get it to realise this is wrong?
+
+        with pytest.raises(ValueError, match="Not a valid grid ufunc signature"):
+
+            @as_grid_ufunc()
+            def ufunc(
+                a: Annotated[np.ndarray, "nonsense"]  # type: ignore
+            ) -> Annotated[np.ndarray, "X:center"]:
+                ...
+
+        with pytest.raises(ValueError, match="Not a valid grid ufunc signature"):
+
+            @as_grid_ufunc()
+            def ufunc(
+                a: Annotated[np.ndarray, "X:Mars"]
+            ) -> Annotated[np.ndarray, "X:center"]:
+                ...
+
+    @pytest.mark.xfail(reason="signature regex will assume nonsense==no inputs")
+    def test_invalid_return_arg_annotation(self):
+        # TODO how to get it to realise this is wrong?
+
+        with pytest.raises(ValueError, match="Not a valid grid ufunc signature"):
+
+            @as_grid_ufunc()
+            def ufunc(
+                a: Annotated[np.ndarray, "X:center"]
+            ) -> Annotated[np.ndarray, "X:Venus"]:
+                ...
+
+    def test_both_sig_kwarg_and_hints_given(self):
+        with pytest.raises(
+            ValueError, match="only one of either type hints or signature kwarg"
+        ):
+
+            @as_grid_ufunc(signature="(X:center)->(X:left)")
+            def ufunc(
+                a: Annotated[np.ndarray, "X:center"]
+            ) -> Annotated[np.ndarray, "X:left"]:
+                ...
+
+    def test_type_hint_as_numpy_ndarray(self):
+
+        # This should raise a mypy error, which is then ignored
+        @as_grid_ufunc()
+        def ufunc1(a: Annotated[str, "X:center"]) -> Annotated[np.ndarray, "X:center"]:
+            # np.ndarray has a .strides method but str doesn't (and nor does xr.DataArray)
+            print(a.strides)  # type: ignore
+            return a  # type: ignore
+
+        # This should pass mypy without raising any errors
+        @as_grid_ufunc()
+        def ufunc3(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:center"]:
+            print(a.strides)
+            return a
 
 
 def create_1d_test_grid_ds(ax_name, length):
@@ -166,8 +300,10 @@ class TestGridUFuncNoPadding:
     def test_stores_ufunc_kwarg_info(self):
         signature = "(X:center)->(X:left)"
 
-        @as_grid_ufunc(signature)
-        def diff_center_to_left(a):
+        @as_grid_ufunc()
+        def diff_center_to_left(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:left"]:
             return a - np.roll(a, shift=-1)
 
         assert isinstance(diff_center_to_left, GridUFunc)
@@ -175,7 +311,7 @@ class TestGridUFuncNoPadding:
 
         with pytest.raises(TypeError, match="Unsupported keyword argument"):
 
-            @as_grid_ufunc(signature, junk="useless")
+            @as_grid_ufunc(junk="useless")
             def diff_center_to_left(a):
                 return a - np.roll(a, shift=-1)
 
@@ -186,12 +322,17 @@ class TestGridUFuncNoPadding:
         grid._ds.drop_vars("depth_o")
         da = np.sin(grid._ds.depth_g * 2 * np.pi / 9)
 
-        with pytest.raises(ValueError, match=re.escape("(depth:outer) does not exist")):
+        with pytest.raises(
+            ValueError,
+            match=re.escape("Axis:positions pair depth:outer does not exist"),
+        ):
+            da: Annotated[np.ndarray, "X:outer"]
             apply_as_grid_ufunc(
                 lambda x: x, da, axis=[("depth",)], grid=grid, signature="(X:outer)->()"
             )
 
         with pytest.raises(ValueError, match="coordinate depth_c does not appear"):
+            da: Annotated[np.ndarray, "X:center"]
             apply_as_grid_ufunc(
                 lambda x: x,
                 da,
@@ -230,8 +371,10 @@ class TestGridUFuncNoPadding:
         assert_equal(result, expected)
 
         # Test decorator
-        @as_grid_ufunc("(X:center)->(X:left)")
-        def diff_center_to_left(a):
+        @as_grid_ufunc()
+        def diff_center_to_left(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:left"]:
             return a - np.roll(a, shift=-1)
 
         result = diff_center_to_left(grid, da, axis=[("depth",)])
@@ -272,8 +415,10 @@ class TestGridUFuncNoPadding:
         assert_equal(result, expected)
 
         # Test decorator
-        @as_grid_ufunc("(X:center)->(X:inner)", dask="parallelized")
-        def interp_center_to_inner(a):
+        @as_grid_ufunc(dask="parallelized")
+        def interp_center_to_inner(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:inner"]:
             return 0.5 * (a[:-1] + a[1:])
 
         result = interp_center_to_inner(grid, da, axis=[("depth",)]).compute()
@@ -319,8 +464,10 @@ class TestGridUFuncNoPadding:
         assert_equal(result, expected)
 
         # Test decorator
-        @as_grid_ufunc("(X:center)->(X:left)", dask="allowed")
-        def diff_overlap(a):
+        @as_grid_ufunc(dask="allowed")
+        def diff_overlap(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:left"]:
             return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
 
         result = diff_overlap(
@@ -356,8 +503,10 @@ class TestGridUFuncNoPadding:
         assert_equal(result, expected)
 
         # Test decorator
-        @as_grid_ufunc("(X:center)->(X:left)")
-        def diff_center_to_left(a):
+        @as_grid_ufunc()
+        def diff_center_to_left(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:left"]:
             return a - np.roll(a, shift=-1, axis=-1)
 
         result = diff_center_to_left(grid, da, axis=[("lon",)])
@@ -397,8 +546,10 @@ class TestGridUFuncNoPadding:
         assert_equal(result, expected)
 
         # Test decorator
-        @as_grid_ufunc("(X:left),(X:right)->()")
-        def inner_product_left_right(a, b):
+        @as_grid_ufunc()
+        def inner_product_left_right(
+            a: Annotated[np.ndarray, "X:left"], b: Annotated[np.ndarray, "X:right"]
+        ):
             return np.inner(a, b)
 
         result = inner_product_left_right(grid, a, b, axis=[("depth",), ("depth",)])
@@ -443,8 +594,13 @@ class TestGridUFuncNoPadding:
         assert_equal(v, expected_v)
 
         # Test decorator
-        @as_grid_ufunc("(X:center,Y:center)->(X:inner,Y:center),(X:center,Y:inner)")
-        def grad_to_inner(a):
+        @as_grid_ufunc()
+        def grad_to_inner(
+            a: Annotated[np.ndarray, "X:center,Y:center"]
+        ) -> Tuple[
+            Annotated[np.ndarray, "X:inner,Y:center"],
+            Annotated[np.ndarray, "X:center,Y:inner"],
+        ]:
             return diff_center_to_inner(a, axis=0), diff_center_to_inner(a, axis=1)
 
         u, v = grad_to_inner(grid, a, axis=[("lon", "lat")])
@@ -665,12 +821,13 @@ class TestDaskOverlap:
 
         # Test decorator
         @as_grid_ufunc(
-            "(X:center)->(X:left)",
             boundary_width={"X": (1, 0)},
             dask="allowed",
             map_overlap=True,
         )
-        def diff_center_to_left(a):
+        def diff_center_to_left(
+            a: Annotated[np.ndarray, "X:center"]
+        ) -> Annotated[np.ndarray, "X:left"]:
             return a[..., 1:] - a[..., :-1]
 
         result = diff_center_to_left(
@@ -726,12 +883,13 @@ class TestDaskOverlap:
 
     def test_raise_when_ufunc_changes_chunksize(self):
         @as_grid_ufunc(
-            "(X:outer)->(X:center)",
             boundary_width={"X": (1, 0)},
             dask="allowed",
             map_overlap=True,
         )
-        def diff_outer_to_center(a):
+        def diff_outer_to_center(
+            a: Annotated[np.ndarray, "X:outer"]
+        ) -> Annotated[np.ndarray, "X:center"]:
             """Mocking up a function which can only act on in-memory arrays, and requires no padding"""
             if isinstance(a, np.ndarray):
                 return a[..., 1:] - a[..., :-1]
@@ -753,12 +911,13 @@ class TestDaskOverlap:
 
     def test_multiple_inputs(self):
         @as_grid_ufunc(
-            "(X:left),(X:right)->(X:center)",
             boundary_width=None,
             map_overlap=True,
             dask="allowed",
         )
-        def multiply_left_right(a, b):
+        def multiply_left_right(
+            a: Annotated[np.ndarray, "X:left"], b: Annotated[np.ndarray, "X:right"]
+        ) -> Annotated[np.ndarray, "X:center"]:
             """Mocking up a function which can only act on in-memory arrays, and requires no padding"""
             if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
                 return np.multiply(a, b)
@@ -845,29 +1004,31 @@ class TestMapOverlapGridops:
 
 class TestSignaturesEquivalent:
     def test_equivalent(self):
-        sig1 = Signature("(X:center)->(X:left)")
-        sig2 = Signature("(X:center)->(X:left)")
+        sig1 = _GridUFuncSignature.from_string("(X:center)->(X:left)")
+        sig2 = _GridUFuncSignature.from_string("(X:center)->(X:left)")
         assert sig1.equivalent(sig2)
 
-        sig3 = Signature("(Y:center)->(Y:left)")
+        sig3 = _GridUFuncSignature.from_string("(Y:center)->(Y:left)")
         assert sig1.equivalent(sig3)
 
     def test_not_equivalent(self):
-        sig1 = Signature("(X:center)->(X:left)")
-        sig2 = Signature("(X:center)->(X:center)")
+        sig1 = _GridUFuncSignature.from_string("(X:center)->(X:left)")
+        sig2 = _GridUFuncSignature.from_string("(X:center)->(X:center)")
         assert not sig1.equivalent(sig2)
 
-        sig3 = Signature("(X:center)->(Y:left)")
+        sig3 = _GridUFuncSignature.from_string("(X:center)->(Y:left)")
         assert not sig1.equivalent(sig3)
 
-        sig4 = Signature("(X:center,X:center)->(X:left)")
+        sig4 = _GridUFuncSignature.from_string("(X:center,X:center)->(X:left)")
         assert not sig1.equivalent(sig4)
 
     def test_no_indices(self):
-        sig = Signature("()->()")
+        sig = _GridUFuncSignature.from_string("()->()")
         assert sig.equivalent(sig)
 
 
+# TODO Is there a way to prevent this running at test module import time?
+# TODO (test by adding a raise in here)
 class GridOpsMockUp:
     """
     Container that stores some mocked-up grid ufuncs to look through.
@@ -889,41 +1050,44 @@ class GridOpsMockUp:
     def diff_center_to_right_extend(a):
         return np.roll(a, 1) - a
 
-    @staticmethod
-    @as_grid_ufunc(signature="()->()")
-    def pass_through_kwargs(**kwargs):
-        return kwargs
-
 
 class TestGridUFuncDispatch:
     def test_select_ufunc(self):
         gridufunc, _ = _select_grid_ufunc(
-            "diff", Signature("(X:center)->(X:left)"), module=GridOpsMockUp
+            "diff",
+            _GridUFuncSignature.from_string("(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", Signature("(Y:center)->(Y:left)"), module=GridOpsMockUp
+            "diff",
+            _GridUFuncSignature.from_string("(Y:center)->(Y:left)"),
+            module=GridOpsMockUp,
         )
         assert gridufunc is GridOpsMockUp.diff_center_to_left
 
         with pytest.raises(NotImplementedError):
             _select_grid_ufunc(
-                "diff", Signature("(X:center)->(Y:left)"), module=GridOpsMockUp
+                "diff",
+                _GridUFuncSignature.from_string("(X:center)->(Y:left)"),
+                module=GridOpsMockUp,
             )
 
     def test_select_ufunc_wrong_signature(self):
         with pytest.raises(NotImplementedError):
             _select_grid_ufunc(
-                "diff", Signature("(X:center)->(X:center)"), module=GridOpsMockUp
+                "diff",
+                _GridUFuncSignature.from_string("(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",
-            Signature("(X:center)->(X:right)"),
+            _GridUFuncSignature.from_string("(X:center)->(X:right)"),
             module=GridOpsMockUp,
             boundary="fill",
         )
@@ -932,7 +1096,7 @@ class TestGridUFuncDispatch:
         with pytest.raises(NotImplementedError):
             _select_grid_ufunc(
                 "diff",
-                Signature("(X:center)->(X:right)"),
+                _GridUFuncSignature.from_string("(X:center)->(X:right)"),
                 module=GridOpsMockUp,
                 boundary="nonsense",
             )
@@ -941,6 +1105,9 @@ class TestGridUFuncDispatch:
     def test_pass_through_other_kwargs(self):
         # TODO put this in test_grid.py instead?
         gridufunc, _ = _select_grid_ufunc(
-            "pass", Signature("()->()"), module=GridOpsMockUp, boundary="fill"
+            "pass",
+            _GridUFuncSignature.from_string("()->()"),
+            module=GridOpsMockUp,
+            boundary="fill",
         )
         assert gridufunc(a=1) == {"a": 1}