diff --git a/src/fmu/dataio/_model/fields.py b/src/fmu/dataio/_model/fields.py index 53b3e699b..ca2ea29d1 100644 --- a/src/fmu/dataio/_model/fields.py +++ b/src/fmu/dataio/_model/fields.py @@ -134,6 +134,14 @@ class Parameters(RootModel): root: Dict[str, Union[Parameters, int, float, str]] """A dictionary representing parameters as-is from parameters.txt.""" + def __iter__(self) -> Any: + # Using ´Any´ as return type here as mypy is having issues + # resolving the correct type + return iter(self.root) + + def __getitem__(self, item: str) -> Union[Parameters, int, float, str]: + return self.root[item] + class Aggregation(BaseModel): """ diff --git a/src/fmu/dataio/_model/global_configuration.py b/src/fmu/dataio/_model/global_configuration.py index f34a69fc1..04b0120b8 100644 --- a/src/fmu/dataio/_model/global_configuration.py +++ b/src/fmu/dataio/_model/global_configuration.py @@ -108,6 +108,14 @@ class Stratigraphy(RootModel[Dict[str, StratigraphyElement]]): A collection of StratigraphyElement instances, accessible by keys. """ + def __iter__(self) -> Any: + # Using ´Any´ as return type here as mypy is having issues + # resolving the correct type + return iter(self.root) + + def __getitem__(self, item: str) -> StratigraphyElement: + return self.root[item] + class GlobalConfiguration(BaseModel): """ diff --git a/tests/test_units/test_dunder_methods.py b/tests/test_units/test_dunder_methods.py new file mode 100644 index 000000000..e3e1cb151 --- /dev/null +++ b/tests/test_units/test_dunder_methods.py @@ -0,0 +1,94 @@ +"""Explicitly test all dunder methods.""" + +import pytest + +from fmu.dataio._model.fields import Parameters +from fmu.dataio._model.global_configuration import Stratigraphy, StratigraphyElement + +# -------------------------------------------------------------------------------------- +# Stratigraphy +# -------------------------------------------------------------------------------------- + + +@pytest.fixture(name="testdata_stratigraphy", scope="function") +def _fixture_testdata_stratigraphy() -> Stratigraphy: + """ + Return a dict of StratigraphyElement instances. + """ + return Stratigraphy( + root={ + "TopStratUnit1": StratigraphyElement( + name="Stratigraphic Unit 1", + stratigraphic=True, + alias=["TopSU1", "TopLayer1"], + ), + "TopStratUnit2": StratigraphyElement( + name="Stratigraphic Unit 2", + stratigraphic=True, + alias=["TopSU2", "TopLayer2"], + ), + "TopStratUnit3": StratigraphyElement( + name="Stratigraphic Unit 3", + stratigraphic=True, + alias=["TopSU3", "TopLayer3"], + ), + } + ) + + +def test_stratigraphy_dunder_iter(testdata_stratigraphy): + try: + count = 0 + for item in testdata_stratigraphy: + count += 1 + assert count == 3 + except Exception: + assert False, "Stratigraphy class does not have __iter__()" + + +def test_stratigraphy_dunder_getitem(testdata_stratigraphy): + try: + testdata_stratigraphy["TopStratUnit2"] + except Exception: + assert False, "Stratigraphy class does not have __getitem__()" + + +# -------------------------------------------------------------------------------------- +# Parameters +# -------------------------------------------------------------------------------------- + + +@pytest.fixture(name="testdata_parameters", scope="function") +def _fixture_testdata_parameters() -> Parameters: + """ + Return a nested dict of Parameter instances. + """ + return Parameters( + root={ + "p1": 42, + "p2": "not so nested", + "p3": Parameters( + root={ + "p3_1": 42.3, + "p3_2": "more nested", + } + ), + } + ) + + +def test_parameters_dunder_iter(testdata_parameters): + try: + count = 0 + for item in testdata_parameters: + count += 1 + assert count == 3 + except Exception: + assert False, "Parameters class does not have __iter__()" + + +def test_parameters_dunder_getitem(testdata_parameters): + try: + testdata_parameters["p2"] + except Exception: + assert False, "Parameters class does not have __getitem__()"