Modify to multi variable constraint function

src/CSET/operators/constraints.py

@@ -26,6 +26,9 @@
 import CSET.operators._utils as operator_utils
 from CSET._common import iter_maybe
 
+# STASH code pattern: mXXsXXiXXX where X is a digit
+_STASH_RE = re.compile(r"^m\d{2}s\d{2}i\d{3}$")
+
 
 def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint:
     """Generate constraint from STASH code.
@@ -52,33 +55,44 @@ def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint:
 def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint:
     """Generate constraint from variable name or STASH code.
 
-    Operator that takes a CF compliant variable name string, and generates an
+    Operator that takes a CF compliant variable name string or list of names, and generates an
     iris constraint to be passed into the read or filter operator. Can also be
     passed a STASH code to generate a STASH constraint.
 
     Arguments
     ---------
-    varname: str
-        CF compliant name of variable, or a UM STASH code such as "m01s03i236".
+    varname: str | list[str]
+        CF compliant name(s) of variable, or a UM STASH code such as "m01s03i236".
 
     Returns
     -------
     varname_constraint: iris.Constraint
+        If a single UM STASHcode is requested, varname constraint is by STASHcode
+        If a single variable name is requested, constraint by varname
+        If multiple variable names are requested, constrain by list of variables.
     """
-    if re.match(r"m[0-9]{2}s[0-9]{2}i[0-9]{3}$", varname):
-        varname_constraint = iris.AttributeConstraint(STASH=varname)
-    else:
-        varname_constraint = iris.Constraint(name=varname)
+    # Case 1: UM STASHcode input
+    if isinstance(varname, str) and _STASH_RE.match(varname):
+        return iris.AttributeConstraint(STASH=varname)
 
     # Ensure access to variable vector components for computed fields
-    if varname == "wind_speed_at_10m":
+    if "wind_speed_at_10m" in iter_maybe(varname):
+        varname = [varname]
+        varname.extend(["eastward_wind_at_10m", "northward_wind_at_10m"])
+
+    # Case 2: Multiple varnames
+    if isinstance(varname, (list, tuple)):
         varname_constraint = iris.Constraint(
             cube_func=lambda cube: (
-                cube.long_name
-                in ["wind_at_10m", "eastward_wind_at_10m", "northward_wind_at_10m"]
+                cube.long_name in varname
+                or cube.standard_name in varname
+                or cube.var_name in varname
             )
         )
 
+    else:
+        varname_constraint = iris.Constraint(name=varname)
+
     return varname_constraint
 
 

tests/operators/test_constraints.py

@@ -16,6 +16,8 @@
 
 from datetime import datetime
 
+import iris
+import numpy as np
 import pytest
 
 from CSET.operators import constraints
@@ -288,6 +290,33 @@ def test_generate_attribute_constraint_with_value():
     assert expected_attr_constraint in repr(attr_constraint)
 
 
+def test_generate_var_constraint_multiple_names():
+    """Test constraint works for multiple variable names."""
+    # Create two cubes with different names
+    cube1 = iris.cube.Cube(np.arange(5), long_name="temperature_long")
+    cube1.var_name = "var_temperature"
+    cube2 = iris.cube.Cube(np.arange(5), standard_name="wind_speed")
+    # Third cube that should NOT match
+    cube3 = iris.cube.Cube(np.arange(5), long_name="surface_pressure")
+    # Generate constraint with multiple names
+    constraint = constraints.generate_var_constraint(["var_temperature", "wind_speed"])
+    # Apply constraint
+    cubes = iris.cube.CubeList([cube1, cube2, cube3])
+    result = cubes.extract(constraint)
+    # Check correct cubes are selected
+    result_names = [c.name() for c in result]
+
+    assert cube1 in result
+    assert cube2 in result
+    assert cube3 not in result
+
+    assert "temperature_long" in result_names
+    assert "wind_speed" in result_names
+    assert "surface_pressure" not in result_names
+    # Should only return 2 cubes
+    assert len(result) == 2
+
+
 def test_generate_remove_single_level_constraint():
     """Tests constraint to remove default model_level_number of zero."""
     remove_level_constraint = constraints.generate_remove_single_level_constraint(