Coverage for python/lsst/daf/butler/queries/tree/_column_reference.py: 63%
74 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-04 02:55 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-04 02:55 -0700
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__ = ("ColumnReference", "DimensionKeyReference", "DimensionFieldReference", "DatasetFieldReference")
32from typing import TYPE_CHECKING, ClassVar, Literal, TypeAlias, TypeVar, Union, final
34import pydantic
36from ...column_spec import ColumnType
37from ...dimensions import Dimension, DimensionElement
38from ._base import ColumnExpressionBase, DatasetFieldName, InvalidQueryError
40if TYPE_CHECKING:
41 from ..visitors import ColumnExpressionVisitor
42 from ._column_set import ColumnSet
45_T = TypeVar("_T")
48@final
49class DimensionKeyReference(ColumnExpressionBase):
50 """A column expression that references a dimension primary key column."""
52 expression_type: Literal["dimension_key"] = "dimension_key"
54 is_column_reference: ClassVar[bool] = True
56 dimension: Dimension
57 """Definition and name of this dimension."""
59 def gather_required_columns(self, columns: ColumnSet) -> None:
60 # Docstring inherited.
61 columns.update_dimensions(self.dimension.minimal_group)
63 @property
64 def column_type(self) -> ColumnType:
65 # Docstring inherited.
66 return self.dimension.primary_key.type
68 def __str__(self) -> str:
69 return self.dimension.name
71 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
72 # Docstring inherited.
73 return visitor.visit_dimension_key_reference(self)
76@final
77class DimensionFieldReference(ColumnExpressionBase):
78 """A column expression that references a dimension record column that is
79 not a primary key.
80 """
82 expression_type: Literal["dimension_field"] = "dimension_field"
84 is_column_reference: ClassVar[bool] = True
86 element: DimensionElement
87 """Definition and name of the dimension element."""
89 field: str
90 """Name of the field (i.e. column) in the element's logical table."""
92 def gather_required_columns(self, columns: ColumnSet) -> None:
93 # Docstring inherited.
94 columns.update_dimensions(self.element.minimal_group)
95 columns.dimension_fields[self.element.name].add(self.field)
97 @property
98 def column_type(self) -> ColumnType:
99 # Docstring inherited.
100 return self.element.schema.remainder[self.field].type
102 def __str__(self) -> str:
103 return f"{self.element}.{self.field}"
105 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
106 # Docstring inherited.
107 return visitor.visit_dimension_field_reference(self)
109 @pydantic.model_validator(mode="after")
110 def _validate_field(self) -> DimensionFieldReference:
111 if self.field not in self.element.schema.remainder.names:
112 raise InvalidQueryError(f"Dimension field {self.element.name}.{self.field} does not exist.")
113 return self
116@final
117class DatasetFieldReference(ColumnExpressionBase):
118 """A column expression that references a column associated with a dataset
119 type.
120 """
122 expression_type: Literal["dataset_field"] = "dataset_field"
124 is_column_reference: ClassVar[bool] = True
126 dataset_type: str
127 """Name of the dataset type."""
129 field: DatasetFieldName
130 """Name of the field (i.e. column) in the dataset's logical table."""
132 def gather_required_columns(self, columns: ColumnSet) -> None:
133 # Docstring inherited.
134 columns.dataset_fields[self.dataset_type].add(self.field)
136 @property
137 def column_type(self) -> ColumnType:
138 # Docstring inherited.
139 match self.field:
140 case "dataset_id":
141 return "uuid"
142 case "ingest_date":
143 return "datetime"
144 case "run":
145 return "string"
146 case "collection":
147 return "string"
148 case "timespan":
149 return "timespan"
150 raise AssertionError(f"Invalid field {self.field!r} for dataset.")
152 def __str__(self) -> str:
153 return f"{self.dataset_type}.{self.field}"
155 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
156 # Docstring inherited.
157 return visitor.visit_dataset_field_reference(self)
160ColumnReference: TypeAlias = Union[
161 DimensionKeyReference,
162 DimensionFieldReference,
163 DatasetFieldReference,
164]