Skip to content

Commit

Permalink
Add Pydantic v1 version of ColumnSpec.
Browse files Browse the repository at this point in the history
It should be safe to revert this commit once we require Pydantic v2.
  • Loading branch information
TallJimbo committed Dec 17, 2023
1 parent 93516ca commit 1dd8413
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 17 deletions.
10 changes: 9 additions & 1 deletion python/lsst/daf/butler/column_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,10 +50,11 @@
from lsst.sphgeom import Region

from . import arrow_utils, ddl
from ._compat import _BaseModelCompat
from ._timespan import Timespan


class ColumnSpec(pydantic.BaseModel, ABC):
class ColumnSpec(_BaseModelCompat, ABC):
"""Base class for descriptions of table columns."""

name: str = pydantic.Field(description="""Name of the column.""")
Expand Down Expand Up @@ -183,6 +184,7 @@ class IntColumnSpec(ColumnSpec):

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_primitive(self.name, pa.uint64(), nullable=self.nullable)


Expand All @@ -203,6 +205,7 @@ def to_sql_spec(self, **kwargs: Any) -> ddl.FieldSpec:

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_primitive(self.name, pa.string(), nullable=self.nullable)


Expand All @@ -223,6 +226,7 @@ def to_sql_spec(self, **kwargs: Any) -> ddl.FieldSpec:

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_primitive(
self.name,
# The size for Arrow binary columns is a fixed size, not a maximum
Expand All @@ -242,6 +246,7 @@ class FloatColumnSpec(ColumnSpec):

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_primitive(self.name, pa.float64(), nullable=self.nullable)


Expand All @@ -255,6 +260,7 @@ class BoolColumnSpec(ColumnSpec):

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_primitive(self.name, pa.bool_(), nullable=self.nullable)

Check warning on line 264 in python/lsst/daf/butler/column_spec.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/column_spec.py#L263-L264

Added lines #L263 - L264 were not covered by tests


Expand All @@ -273,6 +279,7 @@ class RegionColumnSpec(ColumnSpec):

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_region(self.name, nullable=self.nullable)


Expand All @@ -288,6 +295,7 @@ class TimespanColumnSpec(ColumnSpec):

def to_arrow(self) -> arrow_utils.ToArrow:
# Docstring inherited.
assert self.nullable is not None, "nullable=None should be resolved by validators"
return arrow_utils.ToArrow.for_timespan(self.name, nullable=self.nullable)


Expand Down
51 changes: 35 additions & 16 deletions python/lsst/daf/butler/dimensions/_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import pydantic
from lsst.resources import ResourcePath, ResourcePathExpression

from .._compat import PYDANTIC_V2
from .._config import Config, ConfigSubset
from .._topology import TopologicalSpace
from ..column_spec import default_nullable, not_nullable
Expand All @@ -53,13 +54,20 @@
# have a version.
_DEFAULT_NAMESPACE = "daf_butler"

if PYDANTIC_V2:
AnnotatedKeyColumnSpec: TypeAlias = Annotated[
KeyColumnSpec, pydantic.Field(discriminator="type"), pydantic.AfterValidator(not_nullable)
]
AnnotatedMetadataColumnSpec: TypeAlias = Annotated[
MetadataColumnSpec, pydantic.Field(discriminator="type"), pydantic.AfterValidator(default_nullable)
]
else:

AnnotatedKeyColumnSpec: TypeAlias = Annotated[
KeyColumnSpec, pydantic.Field(discriminator="type"), pydantic.AfterValidator(not_nullable)
]
AnnotatedMetadataColumnSpec: TypeAlias = Annotated[
MetadataColumnSpec, pydantic.Field(discriminator="type"), pydantic.AfterValidator(default_nullable)
]
class _KeyColumnModel(pydantic.BaseModel):
spec: KeyColumnSpec = pydantic.Field(discriminator="type")

Check warning on line 67 in python/lsst/daf/butler/dimensions/_config.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/dimensions/_config.py#L66-L67

Added lines #L66 - L67 were not covered by tests

class _MetadataColumnModel(pydantic.BaseModel):
spec: MetadataColumnSpec = pydantic.Field(discriminator="type")

Check warning on line 70 in python/lsst/daf/butler/dimensions/_config.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/dimensions/_config.py#L69-L70

Added lines #L69 - L70 were not covered by tests


class DimensionConfig(ConfigSubset):
Expand Down Expand Up @@ -202,17 +210,28 @@ def _extractElementVisitors(self) -> Iterator[DimensionConstructionVisitor]:
`StandardDimensionCombination` to an under-construction
`DimensionUniverse`.
"""
# MyPy is confused by the typing.Annotated usage and/or how Pydantic
# annotated TypeAdapter.
key_adapter: pydantic.TypeAdapter[KeyColumnSpec] = pydantic.TypeAdapter(
AnnotatedKeyColumnSpec # type: ignore
)
metadata_adapter: pydantic.TypeAdapter[MetadataColumnSpec] = pydantic.TypeAdapter(
AnnotatedMetadataColumnSpec # type: ignore
)
if PYDANTIC_V2:
# MyPy is confused by the typing.Annotated usage and/or how
# Pydantic annotated TypeAdapter.
key_adapter: pydantic.TypeAdapter[KeyColumnSpec] = pydantic.TypeAdapter( # type: ignore
AnnotatedKeyColumnSpec # type: ignore
)
validate_key = key_adapter.validate_python
metadata_adapter: pydantic.TypeAdapter[MetadataColumnSpec] = pydantic.TypeAdapter( # type: ignore
AnnotatedMetadataColumnSpec # type: ignore
)
validate_metadata = metadata_adapter.validate_python
else:

def validate_key(value: Any) -> KeyColumnSpec: # type: ignore[misc]
return not_nullable(_KeyColumnModel(spec=value).spec) # type: ignore[return-value]

Check warning on line 227 in python/lsst/daf/butler/dimensions/_config.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/dimensions/_config.py#L226-L227

Added lines #L226 - L227 were not covered by tests

def validate_metadata(value: Any) -> MetadataColumnSpec: # type: ignore[misc]
return default_nullable(_MetadataColumnModel(spec=value).spec) # type: ignore[return-value]

Check warning on line 230 in python/lsst/daf/butler/dimensions/_config.py

View check run for this annotation

Codecov / codecov/patch

python/lsst/daf/butler/dimensions/_config.py#L229-L230

Added lines #L229 - L230 were not covered by tests

for name, subconfig in self["elements"].items():
metadata_columns = [metadata_adapter.validate_python(c) for c in subconfig.get("metadata", ())]
unique_keys = [key_adapter.validate_python(c) for c in subconfig.get("keys", ())]
metadata_columns = [validate_metadata(c) for c in subconfig.get("metadata", ())]
unique_keys = [validate_key(c) for c in subconfig.get("keys", ())]
if subconfig.get("governor", False):
unsupported = {"required", "implied", "viewOf", "alwaysJoin"}
if not unsupported.isdisjoint(subconfig.keys()):
Expand Down

0 comments on commit 1dd8413

Please sign in to comment.