diff --git a/src/CSET/operators/constraints.py b/src/CSET/operators/constraints.py index 8935cbf6d..28c584fd7 100644 --- a/src/CSET/operators/constraints.py +++ b/src/CSET/operators/constraints.py @@ -14,6 +14,7 @@ """Operators to generate constraints to filter with.""" +import re from collections.abc import Iterable from datetime import datetime @@ -45,22 +46,25 @@ def generate_stash_constraint(stash: str, **kwargs) -> iris.AttributeConstraint: def generate_var_constraint(varname: str, **kwargs) -> iris.Constraint: - """Generate constraint from variable name. + """Generate constraint from variable name or STASH code. - Operator that takes a CF compliant variable name string, and uses iris to - generate a constraint to be passed into the read operator to minimize the - CubeList the read operator loads and speed up loading. + Operator that takes a CF compliant variable name string, 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. Needed later for LFRic. + CF compliant name of variable, or a UM STASH code such as "m01s03i236". Returns ------- varname_constraint: iris.Constraint """ - varname_constraint = iris.Constraint(name=varname) + 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) return varname_constraint diff --git a/tests/operators/test_constraints.py b/tests/operators/test_constraints.py index 74bb88026..b372f9d4a 100644 --- a/tests/operators/test_constraints.py +++ b/tests/operators/test_constraints.py @@ -33,6 +33,13 @@ def test_generate_var_constraint(): assert repr(var_constraint) == expected_var_constraint +def test_generate_var_constraint_stash(): + """Generate iris cube constraint for UM STASH code with var constraint.""" + var_constraint = constraints.generate_var_constraint("m01s03i236") + expected_stash_constraint = "AttributeConstraint({'STASH': 'm01s03i236'})" + assert repr(var_constraint) == expected_stash_constraint + + def test_generate_cell_methods_constraint(): """Generate iris cube constraint for cell methods.""" cell_methods_constraint = constraints.generate_cell_methods_constraint([])