Coverage for python/lsst/daf/butler/_dataset_provenance.py: 11%
179 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-14 01:18 -0800
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__ = ("DatasetProvenance",)
32import re
33import uuid
34from typing import TYPE_CHECKING, Any, Self, TypeAlias
36import pydantic
38from ._dataset_ref import DatasetRef, SerializedDatasetRef
39from .dimensions import DataIdValue
41if TYPE_CHECKING:
42 from collections.abc import Mapping, MutableMapping
44 from ._butler import Butler
46# Types that can be stored in a provenance dictionary.
47_PROV_TYPES: TypeAlias = str | int | float | bool | DataIdValue | uuid.UUID
50class DatasetProvenance(pydantic.BaseModel):
51 """Provenance of a single `DatasetRef`."""
53 inputs: list[SerializedDatasetRef] = pydantic.Field(default_factory=list)
54 """The input datasets."""
55 quantum_id: uuid.UUID | None = None
56 """Identifier of the Quantum that was executed."""
57 extras: dict[uuid.UUID, dict[str, _PROV_TYPES]] = pydantic.Field(default_factory=dict)
58 """Extra provenance information associated with a particular dataset."""
59 _uuids: set[uuid.UUID] = pydantic.PrivateAttr(default_factory=set)
61 @pydantic.model_validator(mode="after")
62 def populate_cache(self) -> Self:
63 for ref in self.inputs:
64 self._uuids.add(ref.id)
65 return self
67 def add_input(self, ref: DatasetRef) -> None:
68 """Add an input dataset to the provenance.
70 Parameters
71 ----------
72 ref : `DatasetRef`
73 A dataset to register as an input.
74 """
75 if ref.id in self._uuids:
76 # Already registered.
77 return
78 self._uuids.add(ref.id)
79 self.inputs.append(ref.to_simple())
81 def add_extra_provenance(self, dataset_id: uuid.UUID, extra: Mapping[str, _PROV_TYPES]) -> None:
82 """Attach extra provenance to a specific dataset.
84 Parameters
85 ----------
86 dataset_id : `uuid.UUID`
87 The ID of the dataset to receive this provenance.
88 extra : `~collections.abc.Mapping` [ `str`, `typing.Any` ]
89 The extra provenance information as a dictionary. The values
90 must be simple Python scalars or scalars that can be serialized
91 by Pydantic and convert to a simple string value.
93 Notes
94 -----
95 The keys in the extra provenance can not include provenance keys of
96 ``run``, ``id``, or ``datasettype`` (in upper or lower case).
97 """
98 if dataset_id not in self._uuids:
99 raise ValueError(f"The given dataset ID {dataset_id} is not known to this provenance instance.")
100 extra_keys = {k.lower() for k in extra.keys()}
101 if overlap := (extra_keys & {"run", "id", "datasettype"}):
102 raise ValueError(f"Extra provenance includes a reserved provenance key: {overlap}")
104 self.extras.setdefault(dataset_id, {}).update(extra)
106 def to_flat_dict(
107 self,
108 ref: DatasetRef | None,
109 /,
110 *,
111 prefix: str = "",
112 sep: str = ".",
113 simple_types: bool = False,
114 use_upper: bool | None = None,
115 ) -> dict[str, _PROV_TYPES]:
116 """Return provenance as a flattened dictionary.
118 Parameters
119 ----------
120 ref : `DatasetRef` or `None`
121 If given, a dataset for which this provenance is relevant and
122 should be included.
123 prefix : `str`, optional
124 A prefix to use for each key in the provenance dictionary.
125 sep : `str`, optional
126 Separator to use to represent hierarchy. Must be a single
127 character. Can not be a number, letter, or underscore (to avoid
128 confusion with provenance keys themselves).
129 simple_types : `bool`, optional
130 If `True` only simple Python types will be used in the returned
131 dictionary, specifically UUIDs will be returned as `str`. If
132 `False`, UUIDs will be returned as `uuid.UUID`. Complex types
133 found in `DatasetProvenance.extras` will be cast to a `str`
134 if `True`.
135 use_upper : `bool` or `None`, optional
136 If `None` the case of the keys matches the case of the first
137 character of the prefix (defined by whether `str.isupper()` returns
138 true, else they will be lower case). If `False` the case will be
139 lower case, and if `True` the case will be upper case.
141 Returns
142 -------
143 prov : `dict`
144 Dictionary representing the provenance. The keys are defined
145 in the notes below.
147 Notes
148 -----
149 Keys from the given dataset (all optional if no dataset is given):
151 :id: UUID of the given dataset.
152 :run: Run of the given dataset.
153 :datasettype: Dataset type of the given dataset.
154 :dataid x: An entry for each required dimension, "x", in the data ID.
156 Each input dataset will have the ``id``, ``run``, and ``datasettype``
157 keys as defined above (but no ``dataid`` key) with an ``input N``
158 prefix where ``N`` starts counting at 0.
160 The quantum ID, if present, will use key ``quantum``.
162 Examples
163 --------
164 >>> provenance.to_flat_dict(
165 ... ref, prefix="lsst.butler", sep=".", simple_types=True
166 ... )
167 {
168 "lsst.butler.id": "ae0fa83d-cc89-41dd-9680-f97ede49f01e",
169 "lsst.butler.run": "test_run",
170 "lsst.butler.datasettype": "data",
171 "lsst.butler.dataid.detector": 10,
172 "lsst.butler.dataid.instrument": "LSSTCam",
173 "lsst.butler.quantum": "d93a735b-08f0-477d-bc95-2cc32d6d898b",
174 "lsst.butler.input.0.id": "3dfd7ba5-5e35-4565-9d87-4b33880ed06c",
175 "lsst.butler.input.0.run": "other_run",
176 "lsst.butler.input.0.datasettype": "astropy_parquet",
177 "lsst.butler.input.1.id": "7a99f6e9-4035-3d68-842e-58ecce1dc935",
178 "lsst.butler.input.1.run": "other_run",
179 "lsst.butler.input.1.datasettype": "astropy_parquet",
180 }
182 Raises
183 ------
184 ValueError
185 Raised if the separator is not a single character.
186 """
187 if len(sep) != 1:
188 raise ValueError(f"Separator for provenance keys must be a single character. Got {sep!r}.")
189 if re.match(r"[_\w\d]$", sep):
190 raise ValueError(
191 f"Separator for provenance keys can not be word character or underscore. Got {sep!r}."
192 )
194 def _make_key(*keys: str | int) -> str:
195 """Make the key in the correct form with simpler API."""
196 return self._make_provenance_key(prefix, sep, use_upper, *keys)
198 prov: dict[str, _PROV_TYPES] = {}
199 if ref is not None:
200 prov[_make_key("id")] = ref.id if not simple_types else str(ref.id)
201 prov[_make_key("run")] = ref.run
202 prov[_make_key("datasettype")] = ref.datasetType.name
203 for k, v in sorted(ref.dataId.required.items()):
204 prov[_make_key("dataid", k)] = v
206 if self.quantum_id is not None:
207 prov[_make_key("quantum")] = self.quantum_id if not simple_types else str(self.quantum_id)
209 for i, input in enumerate(self.inputs):
210 prov[_make_key("input", i, "id")] = input.id if not simple_types else str(input.id)
211 if input.run is not None: # for mypy
212 prov[_make_key("input", i, "run")] = input.run
213 if input.datasetType is not None: # for mypy
214 prov[_make_key("input", i, "datasettype")] = input.datasetType.name
216 if input.id in self.extras:
217 for xk, xv in self.extras[input.id].items():
218 if simple_types and not isinstance(xv, str | float | int | bool):
219 xv = str(xv)
220 prov[_make_key("input", i, xk)] = xv
222 return prov
224 @staticmethod
225 def _make_provenance_key(prefix: str, sep: str, use_upper: bool | None, *keys: str | int) -> str:
226 """Construct provenance key from prefix and separator.
228 Parameters
229 ----------
230 prefix : `str`
231 A prefix to use for each key in the provenance dictionary.
232 sep : `str`
233 Separator to use to represent hierarchy. Must be a single
234 character.
235 use_upper : `bool` or `None`
236 If `True` use upper case for provenance keys, if `False` use lower
237 case, if `None` match the case of the prefix.
238 keys : `tuple` of `str` | `int`
239 Components of key to combine with prefix and separator.
241 Returns
242 -------
243 key : `str`
244 Key to use in dictionary. Case of result will match case of
245 prefix (defaulting to lower case if the first character of
246 prefix has no case).
247 """
248 if use_upper is None:
249 use_upper = prefix[0].isupper() if prefix else False
250 if prefix:
251 prefix += sep
252 k = sep.join(str(kk) for kk in keys)
253 if use_upper:
254 k = k.upper()
255 return f"{prefix}{k}"
257 @staticmethod
258 def _find_prefix_and_sep(prov_dict: Mapping[str, Any]) -> tuple[str, str] | tuple[None, None]:
259 """Given a mapping try to determine the prefix and separator for
260 provenance keys.
262 Parameters
263 ----------
264 prov_dict : `collections.abc.Mapping`
265 Mapping to scan. Assumed to include keys populated by
266 `to_flat_dict`.
268 Returns
269 -------
270 prefix : `str`
271 Prefix given to `to_flat_dict`. `None` if no provenance headers
272 were found.
273 sep : `str`
274 Separator given to `to_flat_dict`. `None` if no provenance headers
275 were found.
277 Raises
278 ------
279 ValueError
280 Raised if more than one value was found for either the prefix or
281 separator.
282 """
283 # Best keys to look for are dataid and input 0.
284 # dataid is only used in ref provenance and always has a separator.
285 # input 0 is always present if there is any input provenance.
287 def _update_matches(match: re.Match, prefixes: set[str], separators: set[str]) -> None:
288 prefix, *seps = match.groups()
289 if prefix:
290 # Will have a separator at the end.
291 prefix, sep = prefix[:-1], prefix[-1]
292 separators.add(sep)
293 prefixes.add(prefix)
294 separators |= set(seps)
296 separators: set[str] = set()
297 prefixes: set[str] = set()
299 # It is possible for there to be no inputs and just a reference
300 # dataset. If that reference dataset has no DATAID then the simple
301 # logic will not work. In that scenario look for the presence
302 # of RUN, DATASETTYPE and ID to spot that provenance exists.
303 # In this case it may not be possible to determine a sep value.
304 backup = {}
306 for k in prov_dict:
307 if match := re.match("(.*)input(.)0(.)(?:id|datasettype|run)$", k, flags=re.IGNORECASE):
308 _update_matches(match, prefixes, separators)
309 elif match := re.match(
310 # Data coordinates are a-z or underscore.
311 "(.*)dataid(.)[a-z_]+$",
312 k,
313 flags=re.IGNORECASE,
314 ):
315 _update_matches(match, prefixes, separators)
316 elif match := re.match(r"(.*)\b(id|datasettype|run)$", k, flags=re.IGNORECASE):
317 prefix, key = match.groups()
318 backup[key.lower()] = prefix
320 if not prefixes:
321 if "run" in backup and "datasettype" in backup and "id" in backup:
322 # Looks like there is a provenance after all. All 3 must be
323 # present.
324 for prefix in backup.values():
325 if prefix:
326 prefix, sep = prefix[:-1], prefix[-1]
327 else:
328 sep = " " # Will not be used.
329 prefixes.add(prefix)
330 separators.add(sep)
332 if not prefixes:
333 return None, None
335 if len(separators) > 1:
336 raise ValueError(
337 f"Inconsistent values found for separators in provenance header. Got {separators}."
338 )
339 if len(prefixes) > 1:
340 raise ValueError(f"Inconsistent values for prefix found in provenance headers. Got {prefixes}.")
341 return prefixes.pop(), separators.pop()
343 @classmethod
344 def _find_provenance_keys_in_flat_dict(cls, prov_dict: Mapping[str, Any]) -> dict[str, str]:
345 """Find the provenance keys in a dictionary.
347 Parameters
348 ----------
349 prov_dict : `collections.abc.Mapping`
350 Dictionary to be analyzed. Assumed to have been populated by
351 `to_flat_dict`.
353 Returns
354 -------
355 prov_keys : `dict` [ `str`, `str` ]
356 Provenance key as found in the given header mapping to the
357 standardized provenance key (with prefix removed and "."
358 separator).
359 """
360 prefix, sep = cls._find_prefix_and_sep(prov_dict)
362 if prefix is None:
363 return {}
364 # for mypy which can not work out that the above method returns
365 # str, str or None, None.
366 if sep is None:
367 return {}
368 if prefix:
369 # Prefix will always include the separator if it is defined.
370 prefix += sep
372 core_provenance = tuple(f"{prefix}{k}".lower() for k in ("run", "id", "datasettype", "quantum"))
374 # Need to escape the prefix and separator for regex usage.
375 esc_sep = re.escape(sep)
376 esc_prefix = re.escape(prefix)
377 prov_keys: dict[str, str] = {}
378 for k in list(prov_dict):
379 # the input provenance can include extra keys that we cannot
380 # know so have to match solely on INPUT N.
381 found_key = False
382 if re.match(rf"{esc_prefix}input{esc_sep}(\d+){esc_sep}(.*)$", k, flags=re.IGNORECASE):
383 found_key = True
384 elif k.lower() in core_provenance:
385 found_key = True
386 elif re.match(f"{esc_prefix}dataid{esc_sep}[a-z_]+$", k, flags=re.IGNORECASE):
387 found_key = True
389 if found_key:
390 standard = k.removeprefix(prefix)
391 standard = standard.replace(sep, ".")
392 prov_keys[k] = standard.lower()
394 return prov_keys
396 @classmethod
397 def strip_provenance_from_flat_dict(cls, prov_dict: MutableMapping[str, Any]) -> None:
398 """Remove provenance keys from a mapping that had been populated
399 by `to_flat_dict`.
401 Parameters
402 ----------
403 prov_dict : `collections.abc.MutableMapping`
404 Dictionary to modify.
405 """
406 for prov_key in cls._find_provenance_keys_in_flat_dict(prov_dict):
407 del prov_dict[prov_key]
409 return
411 @classmethod
412 def from_flat_dict(cls, prov_dict: Mapping[str, Any], butler: Butler) -> tuple[Self, DatasetRef | None]:
413 """Create a provenance object from a provenance dictionary.
415 Parameters
416 ----------
417 prov_dict : `collections.abc.Mapping`
418 Dictionary populated by `to_flat_dict`.
419 butler : `lsst.daf.butler.Butler`
420 Butler to query to find references datasets.
422 Returns
423 -------
424 prov : `DatasetProvenance`
425 Provenance extracted from this object.
426 ref : `DatasetRef` or `None`
427 Dataset associated with this provenance. Can be `None` if no
428 provenance found.
430 Raises
431 ------
432 ValueError
433 Raised if no provenance values are found in the dictionary.
434 RuntimeError
435 Raised if a referenced dataset is not known to the given butler.
436 """
437 prov_keys = cls._find_provenance_keys_in_flat_dict(prov_dict)
438 if not prov_keys:
439 raise ValueError("No provenance information found in header.")
441 def _coerce_id(id_: str | uuid.UUID) -> uuid.UUID:
442 if isinstance(id_, uuid.UUID):
443 return id_
444 return uuid.UUID(hex=id_)
446 quantum_id = None
447 ref_id = None
448 input_ids = {}
449 extras: dict[int, dict[str, Any]] = {}
451 for k, standard in prov_keys.items():
452 if standard == "id":
453 ref_id = _coerce_id(prov_dict[k])
454 elif standard == "quantum":
455 quantum_id = _coerce_id(prov_dict[k])
456 elif match := re.match(r"input.(\d+).([a-z_]+)$", standard):
457 input_num = int(match.group(1))
458 subkey = match.group(2)
459 if subkey == "id":
460 input_ids[input_num] = _coerce_id(prov_dict[k])
461 elif subkey not in ("datasettype", "run"):
462 # Extra information. Can not know the original case
463 # but can match to the original dictionary.
464 if k[0].isupper():
465 subkey = subkey.upper()
466 extras.setdefault(input_num, {})[subkey] = prov_dict[k]
468 ref = None
469 if ref_id is not None:
470 ref = butler.get_dataset(ref_id)
471 if ref is None:
472 raise ValueError(
473 f"Dataset associated with this provenance ({ref_id}) is not known to this butler."
474 )
476 provenance = cls(quantum_id=quantum_id)
478 for i in sorted(input_ids):
479 input_ref = butler.get_dataset(input_ids[i])
480 if input_ref is None:
481 raise ValueError(f"Input dataset ({input_ids[i]}) is not known to this butler.")
482 provenance.add_input(input_ref)
483 if i in extras:
484 provenance.add_extra_provenance(input_ref.id, extras[i])
486 return provenance, ref