Skip to content

Commit

Permalink
Fix mypy validation error by adding explicit type guard function
Browse files Browse the repository at this point in the history
Fixes:

python/lsst/daf/butler/queries/expression_factory.py:468: error: Argument "field" to "DatasetFieldReference" has incompatible type "str";
expected "Literal['dataset_id', 'ingest_date', 'run', 'collection', 'timespan']"  [arg-type]

triggered by pydantic 2.9.0.
  • Loading branch information
timj committed Sep 6, 2024
1 parent 1dbd433 commit 3229f5c
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 2 deletions.
2 changes: 1 addition & 1 deletion python/lsst/daf/butler/queries/expression_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -463,7 +463,7 @@ def __repr__(self) -> str:
# to include Datastore record fields.

def __getattr__(self, field: str) -> ScalarExpressionProxy:
if field not in tree.DATASET_FIELD_NAMES:
if not tree.is_dataset_field(field):
raise AttributeError(field)
expression = tree.DatasetFieldReference(dataset_type=self._dataset_type, field=field)
return ResolvedScalarExpressionProxy(expression)
Expand Down
19 changes: 18 additions & 1 deletion python/lsst/daf/butler/queries/tree/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,11 @@
"ColumnExpressionBase",
"DatasetFieldName",
"DATASET_FIELD_NAMES",
"is_dataset_field",
)

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeVar, cast, get_args
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeGuard, TypeVar, cast, get_args

import pydantic

Expand Down Expand Up @@ -63,6 +64,22 @@
_O = TypeVar("_O")


def is_dataset_field(s: str) -> TypeGuard[DatasetFieldName]:
"""Validate a field name.
Parameters
----------
s : `str`
The field name to test.
Returns
-------
is_field : `bool`
Whether or not this is a dataset field.
"""
return s in DATASET_FIELD_NAMES


class QueryTreeBase(pydantic.BaseModel):
"""Base class for all non-primitive types in a query tree."""

Expand Down

0 comments on commit 3229f5c

Please sign in to comment.