Move TVectorField to _typing module

Author: VeckoTheGecko

src/parcels/_core/fieldset.py

@@ -9,19 +9,18 @@
 import uxarray as ux
 import xarray as xr
 
+import parcels._typing as ptyping
 from parcels._core.field import Field, VectorField
 from parcels._core.model import (
     CONSTANT_FIELD_MODELS,
     ModelData,
     StructuredModelData,
-    TVectorField,
     UnstructuredModelData,
 )
 from parcels._core.utils.string import _assert_str_and_python_varname
 from parcels._core.utils.time import get_datetime_type_calendar
 from parcels._core.utils.time import is_compatible as datetime_is_compatible
 from parcels._python import NOTSET, NotSetType
-from parcels._typing import Mesh
 from parcels.interpolators import (
     XConstantField,
 )
@@ -151,7 +150,7 @@ def add_field(self, field: Field, name: str | None = None):
 
         self.fields[name] = field
 
-    def add_constant_field(self, name: str, value, mesh: Mesh = "spherical"):
+    def add_constant_field(self, name: str, value, mesh: ptyping.Mesh = "spherical"):
         """Wrapper function to add a Field that is constant in space,
            useful e.g. when using constant horizontal diffusivity
 
@@ -212,7 +211,7 @@ def from_ugrid_conventions(
         cls,
         ds: ux.UxDataset,
         mesh: str = "spherical",
-        vector_fields: TVectorField | None | NotSetType = NOTSET,
+        vector_fields: ptyping.VectorFields | None | NotSetType = NOTSET,
     ):
         """Create a FieldSet from a Parcels compliant uxarray.UxDataset.
 
@@ -249,8 +248,8 @@ def from_ugrid_conventions(
     def from_sgrid_conventions(
         cls,
         ds: xr.Dataset,
-        mesh: Mesh | None = None,
-        vector_fields: TVectorField | None | NotSetType = NOTSET,
+        mesh: ptyping.Mesh | None = None,
+        vector_fields: ptyping.VectorFields | None | NotSetType = NOTSET,
     ):  # TODO: Update mesh to be discovered from the dataset metadata
         """Create a FieldSet from a dataset using SGRID convention metadata.
 

src/parcels/_core/model.py

@@ -9,6 +9,7 @@
 import xarray as xr
 
 import parcels._sgrid as sgrid
+import parcels._typing as ptyping
 from parcels._core.basegrid import BaseGrid
 from parcels._core.field import Field, VectorField
 from parcels._core.utils.time import TimeInterval
@@ -34,14 +35,12 @@
 )
 from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator
 
-TVectorField = dict[str, tuple[str, str] | tuple[str, str, str]]
-
 
 class ModelData(ABC):
     data: Any
     grid: BaseGrid
     field_to_interpolator: dict[str, ScalarInterpolator | VectorInterpolator]
-    vector_field_components: TVectorField
+    vector_field_components: ptyping.VectorFields
 
     @abstractmethod
     def construct_fields(self) -> list[Field | VectorField]: ...
@@ -84,7 +83,7 @@ def preprocess_sgrid_model_data(ds: xr.Dataset) -> xr.Dataset:
 
 
 class StructuredModelData(ModelData):
-    def __init__(self, data: xr.Dataset, mesh: Mesh, vector_field_components: TVectorField):
+    def __init__(self, data: xr.Dataset, mesh: Mesh, vector_field_components: ptyping.VectorFields):
         if not isinstance(data, xr.Dataset):
             raise ValueError(f"Expected `data` to be an xarray.Dataset . Got {type(data)}")
 
@@ -133,7 +132,7 @@ def construct_fields(self) -> list[Field | VectorField]:
 
     @classmethod
     def from_sgrid_conventions(
-        cls, ds: xr.Dataset, mesh: Mesh | None, vector_fields: TVectorField | None | NotSetType
+        cls, ds: xr.Dataset, mesh: Mesh | None, vector_fields: ptyping.VectorFields | None | NotSetType
     ) -> Self:
         ds = ds.copy()
         if mesh is None:
@@ -172,15 +171,17 @@ def from_sgrid_conventions(
         return model
 
 
-def resolve_vector_fields(ds: xr.Dataset, vector_fields: TVectorField | None | NotSetType) -> TVectorField:
+def resolve_vector_fields(
+    ds: xr.Dataset, vector_fields: ptyping.VectorFields | None | NotSetType
+) -> ptyping.VectorFields:
     if vector_fields is None:
         return {}
     if vector_fields is NOTSET:  # i.e., the default vectorfield discovery behaviour
         return _default_vector_field_components(list(ds.data_vars))
     return vector_fields
 
 
-def assert_vector_field_components_in_dataset(ds: xr.Dataset, vector_fields: TVectorField) -> None:
+def assert_vector_field_components_in_dataset(ds: xr.Dataset, vector_fields: ptyping.VectorFields) -> None:
     for components in vector_fields.values():
         for c in components:
             if c not in ds.data_vars:
@@ -220,7 +221,7 @@ def assert_vector_field_components_in_dataset(ds: xr.Dataset, vector_fields: TVe
 
 
 class UnstructuredModelData(ModelData):
-    def __init__(self, data: ux.UxDataset, grid: UxGrid, vector_field_components: TVectorField):
+    def __init__(self, data: ux.UxDataset, grid: UxGrid, vector_field_components: ptyping.VectorFields):
         if not isinstance(data, ux.UxDataset):
             raise ValueError(f"Expected `data` to be an uxarray.UxDataset . Got {type(data)}")
 
@@ -259,7 +260,9 @@ def scalar_field_names(self) -> list[str]:
         return list(self.data.data_vars)
 
     @classmethod
-    def from_ugrid_conventions(cls, ds: ux.UxDataset, mesh: Mesh, vector_fields: TVectorField | None | NotSetType):
+    def from_ugrid_conventions(
+        cls, ds: ux.UxDataset, mesh: Mesh, vector_fields: ptyping.VectorFields | None | NotSetType
+    ):
         ds_dims = list(ds.dims)
         if not all(dim in ds_dims for dim in ["time", "zf", "zc"]):
             raise ValueError(
@@ -300,9 +303,9 @@ def _get_mesh_type_from_sgrid_dataset(ds_sgrid: xr.Dataset) -> Mesh:
     return "spherical" if _is_coordinate_in_degrees(ds_sgrid[fpoint_x]) else "flat"
 
 
-def _default_vector_field_components(data_vars: Sequence[Hashable]) -> TVectorField:
+def _default_vector_field_components(data_vars: Sequence[Hashable]) -> ptyping.VectorFields:
     vars = set(data_vars)
-    ret: TVectorField = {}
+    ret: ptyping.VectorFields = {}
 
     if {"U", "V"}.issubset(vars):
         ret["UV"] = ("U", "V")

src/parcels/_typing.py

@@ -47,6 +47,7 @@
 CfAxis = XgcmAxisDirection
 XgcmAxisPosition = Literal["center", "left", "right", "inner", "outer"]
 XgcmAxes = Mapping[XgcmAxisDirection, "xgcm.Axis"]
+VectorFields = dict[str, tuple[str, str] | tuple[str, str, str]]
 
 
 def _is_xarray_object(obj):  # with no imports