Move uxarray to typechecking block

src/parcels/_core/fieldset.py

@@ -7,7 +7,6 @@
 
 import cf_xarray  # noqa: F401
 import numpy as np
-import uxarray as ux
 import xarray as xr
 
 import parcels._typing as ptyping
@@ -29,6 +28,8 @@
 )
 
 if TYPE_CHECKING:
+    import uxarray as ux
+
     from parcels._core.basegrid import BaseGrid
     from parcels._typing import TimeLike
 __all__ = ["FieldSet"]

src/parcels/_core/model.py

@@ -3,10 +3,9 @@
 import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Hashable, Sequence
-from typing import Any, Self
+from typing import TYPE_CHECKING, Any, Self
 
 import cf_xarray  # noqa: F401
-import uxarray as ux
 import xarray as xr
 import zarr
 from dask import is_dask_collection
@@ -37,6 +36,9 @@
 )
 from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator
 
+if TYPE_CHECKING:
+    import uxarray as ux
+
 
 class ModelData(ABC):
     data: Any
@@ -320,6 +322,8 @@ def assert_vector_field_components_in_dataset(ds: xr.Dataset, vector_fields: pty
 
 class UnstructuredModelData(ModelData):
     def __init__(self, data: ux.UxDataset, grid: UxGrid, vector_field_components: ptyping.VectorFields):
+        import uxarray as ux
+
         if not isinstance(data, ux.UxDataset):
             raise ValueError(f"Expected `data` to be an uxarray.UxDataset . Got {type(data)}")
 

src/parcels/_core/uxgrid.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-import uxarray as ux
 from dask import is_dask_collection
 
 from parcels._core.basegrid import BaseGrid
 from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, uxgrid_point_in_cell
 from parcels._core.mesh import SphericalMesh, get_mesh
 
+if TYPE_CHECKING:
+    import uxarray as ux
+
 _UXGRID_AXES = Literal["Z", "FACE"]
 
 
@@ -36,6 +38,8 @@ def __init__(
         mesh : str
             The type of mesh used for the grid. Either "flat" or "spherical".
         """
+        import uxarray as ux
+
         if grid.n_max_face_nodes > 3:
             raise ValueError("Provided ux.grid.Grid must contain only triangular cells (n_max_face_nodes=3)")
         self.uxgrid = grid