Improve validation of vector_fields

Author: VeckoTheGecko

src/parcels/_core/model.py

@@ -161,7 +161,7 @@ def from_sgrid_conventions(
         #     ds["lat"] = ds[node_dimensions[1]]
 
         vector_fields = resolve_vector_fields(ds, vector_fields)
-        assert_vector_field_components_in_dataset(ds, vector_fields)
+        assert_valid_vector_fields(ds, vector_fields)
 
         model = cls(ds, mesh=mesh, vector_field_components=vector_fields)
         model._fields = model.construct_fields()
@@ -181,6 +181,29 @@ def resolve_vector_fields(
     return vector_fields
 
 
+def assert_valid_vector_fields(ds: xr.Dataset, vector_fields: ptyping.VectorFields) -> None:
+    # if not isinstance(vector_fields, dict):
+    #     raise ValueError(f"vector_fields must be a dictionary. Got {type(vector_fields)=!r}.")
+
+    for vfield_name, components in vector_fields.items():
+        if not isinstance(vfield_name, str):
+            raise ValueError(
+                f"Invalid `vector_fields` argument. Vector field name in `vector_fields` should be a string. Got field name {vfield_name!r}."
+            )
+        if not (2 <= len(components) <= 3):
+            raise ValueError(
+                f"Invalid `vector_fields` argument. Vector fields must have either 2 or 3 components. Vector field {vfield_name} has {len(components)} components."
+            )
+        for c in components:
+            if not isinstance(c, str):
+                raise ValueError(
+                    f"Invalid `vector_fields` argument. Component names must be strings. Got component name of value {c!r}."
+                )
+
+    assert_vector_field_components_in_dataset(ds, vector_fields)
+    return
+
+
 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:
@@ -273,7 +296,7 @@ def from_ugrid_conventions(
         ds = _discover_ux_U_and_V(ds)
 
         vector_fields = resolve_vector_fields(ds, vector_fields)
-        assert_vector_field_components_in_dataset(ds, vector_fields)
+        assert_valid_vector_fields(ds, vector_fields)
 
         model = cls(ds, grid, vector_fields)
         model._fields = model.construct_fields()

tests/test_fieldset.py

@@ -1,3 +1,4 @@
+from contextlib import nullcontext
 from datetime import timedelta
 
 import cf_xarray  # noqa: F401
@@ -97,6 +98,34 @@ def test_fieldset_from_structured_generic_datasets(ds):
     assert len(fieldset.gridset) == 1
 
 
+@pytest.mark.parametrize(
+    "vector_fields,ctx",
+    [
+        pytest.param(
+            {"UV": ("U",)},
+            pytest.raises(ValueError, match="must have either 2 or 3 components"),
+            id="single-component",
+        ),
+        pytest.param(
+            {"UV": ("U", "missing")},
+            pytest.raises(ValueError, match="not present in the source dataset"),
+            id="component-not-in-dataset",
+        ),
+        pytest.param(
+            {"UV": ("U", "U", "U", "U")},
+            pytest.raises(ValueError, match="must have either 2 or 3 components"),
+            id="too-many-components",
+        ),
+        pytest.param(None, nullcontext(), id="None"),
+    ],
+)
+def test_fieldset_invalid_vector_fields(vector_fields, ctx):
+    ds = datasets_structured["ds_2d_left"][["U_A_grid", "V_A_grid", "grid"]].rename({"U_A_grid": "U", "V_A_grid": "V"})
+
+    with ctx:
+        FieldSet.from_sgrid_conventions(ds, mesh="flat", vector_fields=vector_fields)
+
+
 def test_fieldset_structured_vectorfield_default():
     ds = datasets_structured["ds_2d_left"][["U_A_grid", "V_A_grid", "grid"]].rename({"U_A_grid": "U", "V_A_grid": "V"})