Skip to content

Commit

Permalink
Use v2 model validator for DatasetRef
Browse files Browse the repository at this point in the history
Now use a single model validator that checks everything because
we are checking multiple keys.
  • Loading branch information
timj committed Aug 4, 2023
1 parent 02b7462 commit e5e8b2d
Showing 1 changed file with 36 additions and 20 deletions.
56 changes: 36 additions & 20 deletions python/lsst/daf/butler/core/datasets/ref.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,9 +35,9 @@
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias, runtime_checkable

from lsst.daf.butler._compat import _BaseModelCompat
from lsst.daf.butler._compat import PYDANTIC_V2, _BaseModelCompat
from lsst.utils.classes import immutable
from pydantic import StrictStr, validator
from pydantic import StrictStr

from ..configSupport import LookupKey
from ..dimensions import DataCoordinate, DimensionGraph, DimensionUniverse, SerializedDataCoordinate
Expand Down Expand Up @@ -179,24 +179,40 @@ class SerializedDatasetRef(_BaseModelCompat):
run: StrictStr | None = None
component: StrictStr | None = None

@validator("dataId")
def _check_dataId(cls, v: Any, values: dict[str, Any]) -> Any: # noqa: N805
if (d := "datasetType") in values and values[d] is None:
raise ValueError("Can not specify 'dataId' without specifying 'datasetType'")
return v

@validator("run")
def _check_run(cls, v: Any, values: dict[str, Any]) -> Any: # noqa: N805
if v and (i := "id") in values and values[i] is None:
raise ValueError("'run' cannot be provided unless 'id' is.")
return v

@validator("component")
def _check_component(cls, v: Any, values: dict[str, Any]) -> Any: # noqa: N805
# Component should not be given if datasetType is given
if v and (d := "datasetType") in values and values[d] is not None:
raise ValueError(f"datasetType ({values[d]}) can not be set if component is given ({v}).")
return v
if PYDANTIC_V2:
from pydantic import model_validator

# Can not use "after" validator since in some cases the validator
# seems to trigger with the datasetType field not yet set.
@model_validator(mode="before")
@classmethod
def check_consistent_parameters(cls, data: Any) -> Any:
has_datasetType = data.get("datasetType") is not None
has_dataId = data.get("dataId") is not None
if has_datasetType is not has_dataId:
raise ValueError("If specifying datasetType or dataId, must specify both.")

Check warning on line 193 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L193

Added line #L193 was not covered by tests

if data.get("component") is not None and has_datasetType:
raise ValueError("datasetType can not be set if component is given.")

Check warning on line 196 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L196

Added line #L196 was not covered by tests
return data

else:
from pydantic import validator

Check warning on line 200 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L200

Added line #L200 was not covered by tests

@validator("dataId")

Check warning on line 202 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L202

Added line #L202 was not covered by tests
def _check_dataId(cls, v: Any, values: dict[str, Any]) -> Any: # noqa: N805
print("---\nGetting dataId ", v)
print("Dataset Type: ", values["datasetType"])

Check warning on line 205 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L204-L205

Added lines #L204 - L205 were not covered by tests
if v and (d := "datasetType") in values and values[d] is None:
raise ValueError("Can not specify 'dataId' without specifying 'datasetType'")
return v

Check warning on line 208 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L207-L208

Added lines #L207 - L208 were not covered by tests

@validator("component")

Check warning on line 210 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L210

Added line #L210 was not covered by tests
def _check_component(cls, v: Any, values: dict[str, Any]) -> Any: # noqa: N805
# Component should not be given if datasetType is given
if v and (d := "datasetType") in values and values[d] is not None:
raise ValueError(f"datasetType ({values[d]}) can not be set if component is given ({v}).")
return v

Check warning on line 215 in python/lsst/daf/butler/core/datasets/ref.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/core/datasets/ref.py#L214-L215

Added lines #L214 - L215 were not covered by tests

@classmethod
def direct(
Expand Down

0 comments on commit e5e8b2d

Please sign in to comment.