Coverage for python / lsst / daf / butler / queries / tree / _base.py: 84%
58 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-26 08:49 +0000
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "ANY_DATASET",
32 "DATASET_FIELD_NAMES",
33 "AnyDatasetFieldName",
34 "AnyDatasetType",
35 "ColumnExpressionBase",
36 "DatasetFieldName",
37 "QueryTreeBase",
38 "is_dataset_field",
39)
41import enum
42from abc import ABC, abstractmethod
43from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias, TypeGuard, TypeVar, cast, get_args
45import pydantic
47from ...column_spec import ColumnType
49if TYPE_CHECKING:
50 from ..visitors import ColumnExpressionVisitor
51 from ._column_literal import ColumnLiteral
52 from ._column_set import ColumnSet
55# Type annotation for string literals that can be used as dataset fields in
56# the public API. The 'collection' and 'run' fields are string collection
57# names. Internal interfaces may define other dataset field strings (e.g.
58# collection primary key values) and hence should use `str` rather than this
59# type.
60DatasetFieldName: TypeAlias = Literal["dataset_id", "ingest_date", "run", "collection", "timespan"]
61InternalDatasetFieldName: TypeAlias = Literal["calib_pkey", "collection_key"]
62AnyDatasetFieldName: TypeAlias = DatasetFieldName | InternalDatasetFieldName
64# Tuple of the strings that can be use as dataset fields in public APIs.
65DATASET_FIELD_NAMES: tuple[DatasetFieldName, ...] = tuple(get_args(DatasetFieldName))
67_T = TypeVar("_T")
68_L = TypeVar("_L")
69_A = TypeVar("_A")
70_O = TypeVar("_O")
73class AnyDatasetType(enum.Enum):
74 """A single-element enum that provides a typed sentinal dataset type name
75 for queries over multiple dataset types.
76 """
78 # Since we often use a type-union of ANY_DATASET and str, it's critical
79 # the enum value (which is what Pydantic uses for serialization) not be
80 # a str.
81 ANY_DATASET = -1
83 def __str__(self) -> str:
84 # This string appears in both human-read error messages and the aliases
85 # used for columns in SQL (with quotes, of course). This is a
86 # convenient value because it sorts before and cannot clash with any
87 # regular alphabetical name, while still being human readable.
88 return "[any]"
90 def __repr__(self) -> str:
91 return "ANY_DATASET"
94ANY_DATASET: AnyDatasetType = AnyDatasetType.ANY_DATASET
97def is_dataset_field(s: str) -> TypeGuard[DatasetFieldName]:
98 """Validate a field name.
100 Parameters
101 ----------
102 s : `str`
103 The field name to test.
105 Returns
106 -------
107 is_field : `bool`
108 Whether or not this is a dataset field.
109 """
110 return s in DATASET_FIELD_NAMES
113class QueryTreeBase(pydantic.BaseModel):
114 """Base class for all non-primitive types in a query tree."""
116 model_config = pydantic.ConfigDict(frozen=True, strict=True)
119class ColumnExpressionBase(QueryTreeBase, ABC):
120 """Base class for objects that represent non-boolean column expressions in
121 a query tree.
123 Notes
124 -----
125 This is a closed hierarchy whose concrete, `~typing.final` derived classes
126 are members of the `ColumnExpression` union. That union should be used in
127 type annotations rather than the technically-open base class.
128 """
130 expression_type: str
131 """String literal corresponding to a concrete expression type."""
133 is_literal: ClassVar[bool] = False
134 """Whether this expression wraps a literal Python value."""
136 is_column_reference: ClassVar[bool] = False
137 """Whether this expression wraps a direct reference to column."""
139 @property
140 @abstractmethod
141 def column_type(self) -> ColumnType:
142 """A string enumeration value representing the type of the column
143 expression.
144 """
145 raise NotImplementedError()
147 def get_literal_value(self) -> Any | None:
148 """Return the literal value wrapped by this expression, or `None` if
149 it is not a literal.
150 """
151 return None
153 @abstractmethod
154 def gather_required_columns(self, columns: ColumnSet) -> None:
155 """Add any columns required to evaluate this expression to the
156 given column set.
158 Parameters
159 ----------
160 columns : `ColumnSet`
161 Set of columns to modify in place.
162 """
163 raise NotImplementedError()
165 @abstractmethod
166 def gather_governors(self, governors: set[str]) -> None:
167 """Add any governor dimensions that need to be fully identified for
168 this column expression to be sound.
170 Parameters
171 ----------
172 governors : `set` [ `str` ]
173 Set of governor dimension names to modify in place.
174 """
175 raise NotImplementedError()
177 @abstractmethod
178 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
179 """Invoke the visitor interface.
181 Parameters
182 ----------
183 visitor : `ColumnExpressionVisitor`
184 Visitor to invoke a method on.
186 Returns
187 -------
188 result : `object`
189 Forwarded result from the visitor.
190 """
191 raise NotImplementedError()
194class ColumnLiteralBase(ColumnExpressionBase):
195 """Base class for objects that represent literal values as column
196 expressions in a query tree.
198 Notes
199 -----
200 This is a closed hierarchy whose concrete, `~typing.final` derived classes
201 are members of the `ColumnLiteral` union. That union should be used in
202 type annotations rather than the technically-open base class. The concrete
203 members of that union are only semi-public; they appear in the serialized
204 form of a column expression tree, but should only be constructed via the
205 `make_column_literal` factory function. All concrete members of the union
206 are also guaranteed to have a read-only ``value`` attribute holding the
207 wrapped literal, but it is unspecified whether that is a regular attribute
208 or a `property` (and hence cannot be type-annotated).
209 """
211 is_literal: ClassVar[bool] = True
212 """Whether this expression wraps a literal Python value."""
214 def get_literal_value(self) -> Any:
215 # Docstring inherited.
216 return cast("ColumnLiteral", self).value
218 def gather_required_columns(self, columns: ColumnSet) -> None:
219 # Docstring inherited.
220 pass
222 def gather_governors(self, governors: set[str]) -> None:
223 # Docstring inherited.
224 pass
226 @property
227 def column_type(self) -> ColumnType:
228 # Docstring inherited.
229 return cast(ColumnType, self.expression_type)
231 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
232 # Docstring inherited
233 return visitor.visit_literal(cast("ColumnLiteral", self))