Coverage for python/lsst/pipe/base/caching_limited_butler.py: 26%
77 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-19 01:44 -0800
1# This file is part of pipe_base.
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__ = ["CachingLimitedButler"]
32import logging
33from collections.abc import Iterable, Set
34from typing import Any
36from lsst.daf.butler import (
37 ButlerMetrics,
38 DatasetId,
39 DatasetProvenance,
40 DatasetRef,
41 DeferredDatasetHandle,
42 DimensionUniverse,
43 LimitedButler,
44 StorageClass,
45)
47from ._dataset_handle import InMemoryDatasetHandle
49_LOG = logging.getLogger(__name__)
52class CachingLimitedButler(LimitedButler):
53 """A `LimitedButler` that caches datasets.
55 A `CachingLimitedButler` caches on both `.put()` and `.get()`, and holds a
56 single instance of the most recently used dataset type for that put/get.
58 The dataset types which will be cached on put/get are controlled via the
59 `cache_on_put` and `cache_on_get` attributes, respectively.
61 By default, copies of the cached items are returned on `get`, so that code
62 is free to operate on data in-place. A `no_copy_on_cache` attribute also
63 exists to tell the `CachingLimitedButler` not to return copies when it is
64 known that the calling code can be trusted not to change values, e.g. when
65 passing calibs to `isrTask`.
67 Parameters
68 ----------
69 wrapped : `LimitedButler`
70 The butler to wrap.
71 cache_on_put : `set` [`str`], optional
72 The dataset types to cache on put.
73 cache_on_get : `set` [`str`], optional
74 The dataset types to cache on get.
75 no_copy_on_cache : `set` [`str`], optional
76 The dataset types for which to not return copies when cached.
77 """
79 def __init__(
80 self,
81 wrapped: LimitedButler,
82 cache_on_put: Set[str] = frozenset(),
83 cache_on_get: Set[str] = frozenset(),
84 no_copy_on_cache: Set[str] = frozenset(),
85 ):
86 self._wrapped = wrapped
87 self.storageClasses = self._wrapped.storageClasses
88 self._cache_on_put = cache_on_put
89 self._cache_on_get = cache_on_get
90 self._cache: dict[str, tuple[DatasetId, InMemoryDatasetHandle]] = {}
91 self._no_copy_on_cache = no_copy_on_cache
93 @property
94 def _metrics(self) -> ButlerMetrics:
95 # Need to always forward from the wrapped metrics object.
96 return self._wrapped._metrics
98 @_metrics.setter
99 def _metrics(self, metrics: ButlerMetrics) -> None:
100 # Allow record_metrics() context manager to override the wrapped
101 # butler.
102 self._wrapped._metrics = metrics
104 def get(
105 self,
106 ref: DatasetRef,
107 /,
108 *,
109 parameters: dict[str, Any] | None = None,
110 storageClass: StorageClass | str | None = None,
111 ) -> Any:
112 if storageClass is None:
113 storageClass = ref.datasetType.storageClass
114 elif isinstance(storageClass, str):
115 storageClass = self.storageClasses.getStorageClass(storageClass)
117 # check if we have this dataset type in the cache
118 if cached := self._cache.get(ref.datasetType.name):
119 dataset_id, handle = cached
120 if dataset_id == ref.id: # if we do, check it's the right object
121 _LOG.debug("Returning cached dataset %s", ref)
122 return handle.get(parameters=parameters, storageClass=storageClass)
124 obj = self._wrapped.get(ref, parameters=parameters, storageClass=storageClass)
125 if ref.datasetType.name in self._cache_on_get and not parameters:
126 handle = InMemoryDatasetHandle(
127 obj,
128 storageClass=storageClass,
129 dataId=ref.dataId,
130 copy=ref.datasetType.name not in self._no_copy_on_cache,
131 )
132 # and not parameters is to make sure we don't cache sub-images etc
133 self._cache[ref.datasetType.name] = (ref.id, handle)
134 _LOG.debug("Cached dataset %s", ref)
135 # make sure copy fires if needed
136 return handle.get()
137 return obj
139 def getDeferred(
140 self,
141 ref: DatasetRef,
142 /,
143 *,
144 parameters: dict[str, Any] | None = None,
145 storageClass: str | StorageClass | None = None,
146 ) -> DeferredDatasetHandle:
147 # note that this does not use the cache at all
148 return self._wrapped.getDeferred(ref, parameters=parameters, storageClass=storageClass)
150 def stored_many(self, refs: Iterable[DatasetRef]) -> dict[DatasetRef, bool]:
151 result = {}
152 unknown_refs = []
153 for ref in refs:
154 if cached := self._cache.get(ref.datasetType.name):
155 dataset_id, _ = cached
156 if dataset_id == ref.id:
157 result[ref] = True
158 continue
159 unknown_refs.append(ref)
161 result.update(self._wrapped.stored_many(unknown_refs))
162 return result
164 def isWriteable(self) -> bool:
165 return self._wrapped.isWriteable()
167 def put(self, obj: Any, ref: DatasetRef, /, *, provenance: DatasetProvenance | None = None) -> DatasetRef:
168 if ref.datasetType.name in self._cache_on_put:
169 self._cache[ref.datasetType.name] = (
170 ref.id,
171 InMemoryDatasetHandle(
172 obj,
173 storageClass=ref.datasetType.storageClass,
174 dataId=ref.dataId,
175 copy=ref.datasetType.name not in self._no_copy_on_cache,
176 ),
177 )
178 _LOG.debug("Cached dataset %s on put", ref)
179 return self._wrapped.put(obj, ref, provenance=provenance)
181 def pruneDatasets(
182 self,
183 refs: Iterable[DatasetRef],
184 *,
185 disassociate: bool = True,
186 unstore: bool = False,
187 tags: Iterable[str] = (),
188 purge: bool = False,
189 ) -> None:
190 refs = list(refs)
191 for ref in refs:
192 if cached := self._cache.get(ref.datasetType.name):
193 dataset_id, _ = cached
194 if dataset_id == ref.id:
195 del self._cache[ref.datasetType.name]
197 return self._wrapped.pruneDatasets(
198 refs, disassociate=disassociate, unstore=unstore, tags=tags, purge=purge
199 )
201 @property
202 def dimensions(self) -> DimensionUniverse:
203 return self._wrapped.dimensions
205 @property
206 def _datastore(self) -> Any:
207 return self._wrapped._datastore
209 @_datastore.setter # demanded by MyPy since we declare it to be an instance attribute in LimitedButler.
210 def _datastore(self, value: Any) -> None:
211 self._wrapped._datastore = value