Coverage for python/lsst/daf/butler/queries/tree/_column_reference.py: 65%
80 statements
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-05 11:36 +0000
« prev ^ index » next coverage.py v7.4.3, created at 2024-03-05 11:36 +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__ = ("ColumnReference", "DimensionKeyReference", "DimensionFieldReference", "DatasetFieldReference")
32from typing import TYPE_CHECKING, 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 dimension: Dimension
55 """Definition and name of this dimension."""
57 def gather_required_columns(self, columns: ColumnSet) -> None:
58 # Docstring inherited.
59 columns.update_dimensions(self.dimension.minimal_group)
61 @property
62 def precedence(self) -> int:
63 # Docstring inherited.
64 return 0
66 @property
67 def column_type(self) -> ColumnType:
68 # Docstring inherited.
69 return self.dimension.primary_key.type
71 def __str__(self) -> str:
72 return self.dimension.name
74 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
75 # Docstring inherited.
76 return visitor.visit_dimension_key_reference(self)
79@final
80class DimensionFieldReference(ColumnExpressionBase):
81 """A column expression that references a dimension record column that is
82 not a primary key.
83 """
85 expression_type: Literal["dimension_field"] = "dimension_field"
87 element: DimensionElement
88 """Definition and name of the dimension element."""
90 field: str
91 """Name of the field (i.e. column) in the element's logical table."""
93 def gather_required_columns(self, columns: ColumnSet) -> None:
94 # Docstring inherited.
95 columns.update_dimensions(self.element.minimal_group)
96 columns.dimension_fields[self.element.name].add(self.field)
98 @property
99 def precedence(self) -> int:
100 # Docstring inherited.
101 return 0
103 @property
104 def column_type(self) -> ColumnType:
105 # Docstring inherited.
106 return self.element.schema.remainder[self.field].type
108 def __str__(self) -> str:
109 return f"{self.element}.{self.field}"
111 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
112 # Docstring inherited.
113 return visitor.visit_dimension_field_reference(self)
115 @pydantic.model_validator(mode="after")
116 def _validate_field(self) -> DimensionFieldReference:
117 if self.field not in self.element.schema.remainder.names:
118 raise InvalidQueryError(f"Dimension field {self.element.name}.{self.field} does not exist.")
119 return self
122@final
123class DatasetFieldReference(ColumnExpressionBase):
124 """A column expression that references a column associated with a dataset
125 type.
126 """
128 expression_type: Literal["dataset_field"] = "dataset_field"
130 dataset_type: str
131 """Name of the dataset type."""
133 field: DatasetFieldName
134 """Name of the field (i.e. column) in the dataset's logical table."""
136 def gather_required_columns(self, columns: ColumnSet) -> None:
137 # Docstring inherited.
138 columns.dataset_fields[self.dataset_type].add(self.field)
140 @property
141 def precedence(self) -> int:
142 # Docstring inherited.
143 return 0
145 @property
146 def column_type(self) -> ColumnType:
147 # Docstring inherited.
148 match self.field:
149 case "dataset_id":
150 return "uuid"
151 case "ingest_date":
152 return "datetime"
153 case "run":
154 return "string"
155 case "collection":
156 return "string"
157 case "timespan":
158 return "timespan"
159 raise AssertionError(f"Invalid field {self.field!r} for dataset.")
161 def __str__(self) -> str:
162 return f"{self.dataset_type}.{self.field}"
164 def visit(self, visitor: ColumnExpressionVisitor[_T]) -> _T:
165 # Docstring inherited.
166 return visitor.visit_dataset_field_reference(self)
169ColumnReference: TypeAlias = Union[
170 DimensionKeyReference,
171 DimensionFieldReference,
172 DatasetFieldReference,
173]