Merge branch 'v4-dev' into reuse-cache-file

Author: VeckoTheGecko

src/parcels/_core/field.py

@@ -18,6 +18,7 @@
     AllParcelsErrorCodes,
     StatusCode,
 )
+from parcels._core.utils.string import _assert_str_and_python_varname
 from parcels._core.utils.time import TimeInterval
 from parcels._core.uxgrid import UxGrid
 from parcels._core.xgrid import XGrid, _transpose_xfield_data_to_tzyx
@@ -101,8 +102,9 @@ def __init__(
             raise ValueError(
                 f"Expected `data` to be a uxarray.UxDataArray or xarray.DataArray object, got {type(data)}."
             )
-        if not isinstance(name, str):
-            raise ValueError(f"Expected `name` to be a string, got {type(name)}.")
+
+        _assert_str_and_python_varname(name)
+
         if not isinstance(grid, (UxGrid, XGrid)):
             raise ValueError(f"Expected `grid` to be a parcels UxGrid, or parcels XGrid object, got {type(grid)}.")
 
@@ -246,6 +248,8 @@ class VectorField:
     def __init__(
         self, name: str, U: Field, V: Field, W: Field | None = None, vector_interp_method: Callable | None = None
     ):
+        _assert_str_and_python_varname(name)
+
         self.name = name
         self.U = U
         self.V = V

src/parcels/_core/fieldset.py

@@ -11,6 +11,7 @@
 
 from parcels._core.converters import Geographic, GeographicPolar
 from parcels._core.field import Field, VectorField
+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._core.xgrid import _DEFAULT_XGCM_KWARGS, XGrid
@@ -163,6 +164,8 @@ def add_constant(self, name, value):
         `Diffusion <../examples/tutorial_diffusion.ipynb>`__
         `Periodic boundaries <../examples/tutorial_periodic_boundaries.ipynb>`__
         """
+        _assert_str_and_python_varname(name)
+
         if name in self.constants:
             raise ValueError(f"FieldSet already has a constant with name '{name}'")
         if not isinstance(value, (float, np.floating, int, np.integer)):

src/parcels/_core/particle.py

@@ -2,13 +2,13 @@
 
 import enum
 import operator
-from keyword import iskeyword
 from typing import Literal
 
 import numpy as np
 
 from parcels._compat import _attrgetter_helper
 from parcels._core.statuscodes import StatusCode
+from parcels._core.utils.string import _assert_str_and_python_varname
 from parcels._core.utils.time import TimeInterval
 from parcels._reprs import _format_list_items_multiline
 
@@ -45,9 +45,7 @@ def __init__(
         to_write: bool | Literal["once"] = True,
         attrs: dict | None = None,
     ):
-        if not isinstance(name, str):
-            raise TypeError(f"Variable name must be a string. Got {name=!r}")
-        _assert_valid_python_varname(name)
+        _assert_str_and_python_varname(name)
 
         try:
             dtype = np.dtype(dtype)
@@ -153,12 +151,6 @@ def _assert_no_duplicate_variable_names(*, existing_vars: list[Variable], new_va
             raise ValueError(f"Variable name already exists: {var.name}")
 
 
-def _assert_valid_python_varname(name):
-    if name.isidentifier() and not iskeyword(name):
-        return
-    raise ValueError(f"Particle variable has to be a valid Python variable name. Got {name=!r}")
-
-
 def get_default_particle(spatial_dtype: np.float32 | np.float64) -> ParticleClass:
     if spatial_dtype not in [np.float32, np.float64]:
         raise ValueError(f"spatial_dtype must be np.float32 or np.float64. Got {spatial_dtype=!r}")

src/parcels/_core/utils/string.py

@@ -0,0 +1,16 @@
+from keyword import iskeyword, kwlist
+
+
+def _assert_str_and_python_varname(name):
+    if not isinstance(name, str):
+        raise TypeError(f"Expected a string for variable name, got {type(name).__name__} instead.")
+
+    msg = f"Received invalid Python variable name {name!r}: "
+
+    if not name.isidentifier():
+        msg += "not a valid identifier. HINT: avoid using spaces, special characters, and starting with a number."
+        raise ValueError(msg)
+
+    if iskeyword(name):
+        msg += f"it is a reserved keyword. HINT: avoid using the following names: {', '.join(kwlist)}"
+        raise ValueError(msg)

src/parcels/_datasets/structured/generic.py

@@ -23,10 +23,10 @@ def _rotated_curvilinear_grid():
         {
             "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
             "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
-            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
+            "U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
         },
         coords={
             "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -92,16 +92,19 @@ def _unrolled_cone_curvilinear_grid():
         new_lon_lat.append((lon + pivot[0], lat + pivot[1]))
 
     new_lon, new_lat = zip(*new_lon_lat, strict=True)
-    LON, LAT = np.array(new_lon).reshape(LON.shape), np.array(new_lat).reshape(LAT.shape)
+    LON, LAT = (
+        np.array(new_lon).reshape(LON.shape),
+        np.array(new_lat).reshape(LAT.shape),
+    )
 
     return xr.Dataset(
         {
             "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
             "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
-            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
+            "U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
         },
         coords={
             "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
@@ -140,10 +143,10 @@ def _unrolled_cone_curvilinear_grid():
         {
             "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
             "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
-            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
+            "U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
         },
         coords={
             "XG": (
@@ -182,10 +185,10 @@ def _unrolled_cone_curvilinear_grid():
         {
             "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
             "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, Z, Y, X)),
-            "U (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (A grid)": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
-            "U (C grid)": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
-            "V (C grid)": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
+            "U_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_A_grid": (["time", "ZG", "YG", "XG"], np.random.rand(T, Z, Y, X)),
+            "U_C_grid": (["time", "ZG", "YC", "XG"], np.random.rand(T, Z, Y, X)),
+            "V_C_grid": (["time", "ZG", "YG", "XC"], np.random.rand(T, Z, Y, X)),
         },
         coords={
             "XG": (

tests/test_field.py

@@ -15,10 +15,27 @@
 def test_field_init_param_types():
     data = datasets_structured["ds_2d_left"]
     grid = XGrid.from_dataset(data)
-    with pytest.raises(ValueError, match="Expected `name` to be a string"):
+
+    with pytest.raises(TypeError, match="Expected a string for variable name, got int instead."):
         Field(name=123, data=data["data_g"], grid=grid)
 
-    with pytest.raises(ValueError, match="Expected `data` to be a uxarray.UxDataArray or xarray.DataArray"):
+    for name in ["a b", "123"]:
+        with pytest.raises(
+            ValueError,
+            match=r"Received invalid Python variable name.*: not a valid identifier. HINT: avoid using spaces, special characters, and starting with a number.",
+        ):
+            Field(name=name, data=data["data_g"], grid=grid)
+
+    with pytest.raises(
+        ValueError,
+        match=r"Received invalid Python variable name.*: it is a reserved keyword. HINT: avoid using the following names:.*",
+    ):
+        Field(name="while", data=data["data_g"], grid=grid)
+
+    with pytest.raises(
+        ValueError,
+        match="Expected `data` to be a uxarray.UxDataArray or xarray.DataArray",
+    ):
         Field(name="test", data=123, grid=grid)
 
     with pytest.raises(ValueError, match="Expected `grid` to be a parcels UxGrid, or parcels XGrid"):
@@ -28,7 +45,11 @@ def test_field_init_param_types():
 @pytest.mark.parametrize(
     "data,grid",
     [
-        pytest.param(ux.UxDataArray(), XGrid.from_dataset(datasets_structured["ds_2d_left"]), id="uxdata-grid"),
+        pytest.param(
+            ux.UxDataArray(),
+            XGrid.from_dataset(datasets_structured["ds_2d_left"]),
+            id="uxdata-grid",
+        ),
         pytest.param(
             xr.DataArray(),
             UxGrid(
@@ -76,7 +97,11 @@ def test_field_init_fail_on_float_time_dim():
     (users are expected to use timedelta64 or datetime).
     """
     ds = datasets_structured["ds_2d_left"].copy()
-    ds["time"] = (ds["time"].dims, np.arange(0, T_structured, dtype="float64"), ds["time"].attrs)
+    ds["time"] = (
+        ds["time"].dims,
+        np.arange(0, T_structured, dtype="float64"),
+        ds["time"].attrs,
+    )
 
     data = ds["data_g"]
     grid = XGrid.from_dataset(ds)
@@ -122,7 +147,12 @@ def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, inval
 
     # Test invalid interpolator with wrong signature
     with pytest.raises(ValueError, match=".*incorrect name.*"):
-        Field(name="test", data=ds["data_g"], grid=grid, interp_method=invalid_interpolator_wrong_signature)
+        Field(
+            name="test",
+            data=ds["data_g"],
+            grid=grid,
+            interp_method=invalid_interpolator_wrong_signature,
+        )
 
 
 def test_vectorfield_invalid_interpolator():
@@ -138,7 +168,12 @@ def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, apply
 
     # Test invalid interpolator with wrong signature
     with pytest.raises(ValueError, match=".*incorrect name.*"):
-        VectorField(name="UV", U=U, V=V, vector_interp_method=invalid_interpolator_wrong_signature)
+        VectorField(
+            name="UV",
+            U=U,
+            V=V,
+            vector_interp_method=invalid_interpolator_wrong_signature,
+        )
 
 
 def test_field_unstructured_z_linear():
@@ -161,18 +196,34 @@ def test_field_unstructured_z_linear():
     P = Field(name="p", data=ds.p, grid=grid, interp_method=UXPiecewiseConstantFace)
 
     # Test above first cell center - for piecewise constant, should return the depth of the first cell center
-    assert np.isclose(P.eval(time=ds.time[0].values, z=[10.0], y=[30.0], x=[30.0], applyConversion=False), 55.555557)
+    assert np.isclose(
+        P.eval(time=ds.time[0].values, z=[10.0], y=[30.0], x=[30.0], applyConversion=False),
+        55.555557,
+    )
     # Test below first cell center, but in the first layer  - for piecewise constant, should return the depth of the first cell center
-    assert np.isclose(P.eval(time=ds.time[0].values, z=[65.0], y=[30.0], x=[30.0], applyConversion=False), 55.555557)
+    assert np.isclose(
+        P.eval(time=ds.time[0].values, z=[65.0], y=[30.0], x=[30.0], applyConversion=False),
+        55.555557,
+    )
     # Test bottom layer  - for piecewise constant, should return the depth of the of the bottom layer cell center
     assert np.isclose(
-        P.eval(time=ds.time[0].values, z=[900.0], y=[30.0], x=[30.0], applyConversion=False), 944.44445801
+        P.eval(time=ds.time[0].values, z=[900.0], y=[30.0], x=[30.0], applyConversion=False),
+        944.44445801,
     )
 
     W = Field(name="W", data=ds.W, grid=grid, interp_method=UXPiecewiseLinearNode)
-    assert np.isclose(W.eval(time=ds.time[0].values, z=[10.0], y=[30.0], x=[30.0], applyConversion=False), 10.0)
-    assert np.isclose(W.eval(time=ds.time[0].values, z=[65.0], y=[30.0], x=[30.0], applyConversion=False), 65.0)
-    assert np.isclose(W.eval(time=ds.time[0].values, z=[900.0], y=[30.0], x=[30.0], applyConversion=False), 900.0)
+    assert np.isclose(
+        W.eval(time=ds.time[0].values, z=[10.0], y=[30.0], x=[30.0], applyConversion=False),
+        10.0,
+    )
+    assert np.isclose(
+        W.eval(time=ds.time[0].values, z=[65.0], y=[30.0], x=[30.0], applyConversion=False),
+        65.0,
+    )
+    assert np.isclose(
+        W.eval(time=ds.time[0].values, z=[900.0], y=[30.0], x=[30.0], applyConversion=False),
+        900.0,
+    )
 
 
 def test_field_constant_in_time():
@@ -185,7 +236,13 @@ def test_field_constant_in_time():
     # Assert that the field can be evaluated at any time, and returns the same value
     time = np.datetime64("2000-01-01T00:00:00")
     P1 = P.eval(time=time, z=[10.0], y=[30.0], x=[30.0], applyConversion=False)
-    P2 = P.eval(time=time + np.timedelta64(1, "D"), z=[10.0], y=[30.0], x=[30.0], applyConversion=False)
+    P2 = P.eval(
+        time=time + np.timedelta64(1, "D"),
+        z=[10.0],
+        y=[30.0],
+        x=[30.0],
+        applyConversion=False,
+    )
     assert np.isclose(P1, P2)
 
 

tests/test_fieldset.py

@@ -20,8 +20,8 @@
 def fieldset() -> FieldSet:
     """Fixture to create a FieldSet object for testing."""
     grid = XGrid.from_dataset(ds, mesh="flat")
-    U = Field("U", ds["U (A grid)"], grid)
-    V = Field("V", ds["V (A grid)"], grid)
+    U = Field("U", ds["U_A_grid"], grid)
+    V = Field("V", ds["V_A_grid"], grid)
     UV = VectorField("UV", U, V)
 
     return FieldSet(
@@ -39,6 +39,17 @@ def test_fieldset_add_constant(fieldset):
     assert fieldset.test_constant == 1.0
 
 
+def test_fieldset_add_constant_int_name(fieldset):
+    with pytest.raises(TypeError, match="Expected a string for variable name, got int instead."):
+        fieldset.add_constant(123, 1.0)
+
+
+@pytest.mark.parametrize("name", ["a b", "123", "while"])
+def test_fieldset_add_constant_invalid_name(fieldset, name):
+    with pytest.raises(ValueError, match=r"Received invalid Python variable name.*"):
+        fieldset.add_constant(name, 1.0)
+
+
 def test_fieldset_add_constant_field(fieldset):
     fieldset.add_constant_field("test_constant_field", 1.0)
 
@@ -54,7 +65,7 @@ def test_fieldset_add_constant_field(fieldset):
 
 def test_fieldset_add_field(fieldset):
     grid = XGrid.from_dataset(ds, mesh="flat")
-    field = Field("test_field", ds["U (A grid)"], grid)
+    field = Field("test_field", ds["U_A_grid"], grid)
     fieldset.add_field(field)
     assert fieldset.test_field == field
 
@@ -67,7 +78,7 @@ def test_fieldset_add_field_wrong_type(fieldset):
 
 def test_fieldset_add_field_already_exists(fieldset):
     grid = XGrid.from_dataset(ds, mesh="flat")
-    field = Field("test_field", ds["U (A grid)"], grid)
+    field = Field("test_field", ds["U_A_grid"], grid)
     fieldset.add_field(field, "test_field")
     with pytest.raises(ValueError, match="FieldSet already has a Field with name 'test_field'"):
         fieldset.add_field(field, "test_field")
@@ -104,12 +115,12 @@ def test_fieldset_gridset_multiple_grids(): ...
 
 def test_fieldset_time_interval():
     grid1 = XGrid.from_dataset(ds, mesh="flat")
-    field1 = Field("field1", ds["U (A grid)"], grid1)
+    field1 = Field("field1", ds["U_A_grid"], grid1)
 
     ds2 = ds.copy()
     ds2["time"] = (ds2["time"].dims, ds2["time"].data + np.timedelta64(timedelta(days=1)), ds2["time"].attrs)
     grid2 = XGrid.from_dataset(ds2, mesh="flat")
-    field2 = Field("field2", ds2["U (A grid)"], grid2)
+    field2 = Field("field2", ds2["U_A_grid"], grid2)
 
     fieldset = FieldSet([field1, field2])
     fieldset.add_constant_field("constant_field", 1.0)
@@ -135,8 +146,8 @@ def test_fieldset_init_incompatible_calendars():
     )
 
     grid = XGrid.from_dataset(ds1, mesh="flat")
-    U = Field("U", ds1["U (A grid)"], grid)
-    V = Field("V", ds1["V (A grid)"], grid)
+    U = Field("U", ds1["U_A_grid"], grid)
+    V = Field("V", ds1["V_A_grid"], grid)
     UV = VectorField("UV", U, V)
 
     ds2 = ds.copy()

tests/test_kernel.py

@@ -18,8 +18,8 @@
 def fieldset() -> FieldSet:
     ds = datasets_structured["ds_2d_left"]
     grid = XGrid.from_dataset(ds, mesh="flat")
-    U = Field("U", ds["U (A grid)"], grid)
-    V = Field("V", ds["V (A grid)"], grid)
+    U = Field("U", ds["U_A_grid"], grid)
+    V = Field("V", ds["V_A_grid"], grid)
     return FieldSet([U, V])