Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: Add dunder accessors to Pydantic RootModels #767

Merged
merged 1 commit into from
Sep 2, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/fmu/dataio/_model/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
8 changes: 8 additions & 0 deletions src/fmu/dataio/_model/global_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
94 changes: 94 additions & 0 deletions tests/test_units/test_dunder_methods.py
Original file line number Diff line number Diff line change
@@ -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__()"