Coverage for python/lsst/daf/butler/registry/queries/_query_common.py: 39%
80 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-10 01:29 -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
30import dataclasses
31from abc import abstractmethod
32from collections.abc import Iterable, Iterator
33from contextlib import contextmanager
34from typing import Any, Generic, Self, TypeVar
36from ..._butler import Butler
37from ...dimensions import DataId
38from ...queries import Query, QueryResultsBase
39from .._registry import CollectionArgType
40from ..wildcards import CollectionWildcard
41from ._results import QueryResultsBase as LegacyQueryResultsBase
44@dataclasses.dataclass(frozen=True)
45class CommonQueryArguments:
46 """Simplified version of the arguments passed to many ``Registry.query*``
47 methods.
48 """
50 dataId: DataId | None
51 dataset_types: list[str]
52 collections: list[str] | None
53 where: str
54 bind: dict[str, Any] | None
55 kwargs: dict[str, int | str]
56 check: bool
58 def replaceCollections(self, collections: list[str]) -> CommonQueryArguments:
59 return dataclasses.replace(self, collections=collections)
62_T = TypeVar("_T", bound=QueryResultsBase)
63_U = TypeVar("_U", bound=QueryResultsBase)
66class LegacyQueryResultsMixin(Generic[_T], LegacyQueryResultsBase):
67 """Implements common methods for the various ``QueryResults`` classes in
68 the legacy query system by forwarding to the new query system.
70 Parameters
71 ----------
72 butler : `Butler`
73 Butler object used to execute queries.
74 args : `CommonQueryArguments`
75 User-facing arguments forwarded from the original ``registry.query*``
76 method.
77 """
79 def __init__(
80 self,
81 butler: Butler,
82 args: CommonQueryArguments,
83 ) -> None:
84 self._butler = butler
85 self._args = args
86 self._limit: int | None = None
87 self._order_by: list[str] = []
89 def count(self, *, exact: bool = True, discard: bool = False) -> int:
90 with self._build_query() as result:
91 return result.count(exact=exact, discard=discard)
93 def any(self, *, execute: bool = True, exact: bool = True) -> bool:
94 with self._build_query() as result:
95 return result.any(execute=execute, exact=exact)
97 def order_by(self, *args: str) -> Self:
98 self._order_by.extend(args)
99 return self
101 def limit(self, limit: int) -> Self:
102 self._limit = limit
103 return self
105 def explain_no_results(self, execute: bool = True) -> Iterable[str]:
106 with self._build_query() as result:
107 return result.explain_no_results(execute=execute)
109 @contextmanager
110 def _build_query(self) -> Iterator[_T]:
111 with self._butler.query() as query:
112 a = self._args
113 for dataset_type in a.dataset_types:
114 query = query.join_dataset_search(dataset_type, a.collections)
115 if not a.check:
116 query = query._skip_governor_validation()
118 result = self._build_result(query)
119 result = self._apply_result_modifiers(result)
120 yield result
122 def _apply_result_modifiers(self, result: _U) -> _U:
123 a = self._args
125 result = result.limit(self._limit)
126 if self._order_by:
127 result = result.order_by(*self._order_by)
129 if a.where:
130 result = result.where(a.where, bind=a.bind)
131 if a.dataId or a.kwargs:
132 id_list = [a.dataId] if a.dataId else []
133 # dataId and kwargs have to be sent together as part of the
134 # same call to where() so that the kwargs can override values
135 # in the data ID.
136 result = result.where(*id_list, **a.kwargs, bind=None)
138 return result
140 @abstractmethod
141 def _build_result(self, query: Query) -> _T:
142 raise NotImplementedError("Subclasses must implement _build_result")
145def resolve_collections(
146 butler: Butler, collections: CollectionArgType | None, doomed_by: list[str] | None = None
147) -> list[str]:
148 """Convert the collection argument used throughout the registry query
149 methods to a concrete list of collections.
151 Parameters
152 ----------
153 butler : `Butler`
154 Butler object used to execute queries.
155 collections : Any
156 Any of the values that can be passed as the collections argument to
157 Registry query methods.
158 doomed_by : `list` [ `str` ]
159 Diagnostic messages will be appended to this list.
161 Returns
162 -------
163 collections : `list` [ `str` ]
164 Concrete list of collection names to be used by the query.
165 """
166 if collections is None:
167 return list(butler.collections.defaults)
169 wildcard = CollectionWildcard.from_expression(collections)
170 if wildcard.patterns:
171 result = list(butler.registry.queryCollections(collections))
172 if not result and doomed_by is not None:
173 doomed_by.append(f"No collections found matching expression {wildcard}")
174 return result
175 else:
176 return list(wildcard.strings)