Coverage for python/lsst/pipe/base/prerequisite_helpers.py: 33%
202 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-23 10:54 +0000
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-23 10:54 +0000
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/>.
28"""Helper classes for finding prerequisite input datasets during
29QuantumGraph generation.
30"""
32from __future__ import annotations
34__all__ = (
35 "SkyPixBoundsBuilder",
36 "PrerequisiteFinder",
37 "PrerequisiteBounds",
38 "TimespanBuilder",
39 "PrerequisiteInfo",
40)
42import dataclasses
43from abc import ABC, abstractmethod
44from collections.abc import Callable, Iterable, Mapping, Sequence
45from typing import cast
47from lsst.daf.butler import (
48 Butler,
49 DataCoordinate,
50 DatasetRef,
51 DatasetType,
52 DimensionElement,
53 Registry,
54 SkyPixDimension,
55 Timespan,
56)
57from lsst.daf.butler.registry import MissingDatasetTypeError
58from lsst.sphgeom import RangeSet, Region
60from .pipeline_graph import DatasetTypeNode, PipelineGraph, ReadEdge, TaskNode
63@dataclasses.dataclass
64class PrerequisiteInfo:
65 """A QuantumGraph-generation helper class that manages the searches for all
66 prerequisite input connections for a task.
68 Parameters
69 ----------
70 task_node : `TaskNode`
71 The relevant node.
72 pipeline_graph : `PipelineGraph`
73 The pipeline graph.
74 """
76 bounds: PrerequisiteBounds
77 """Another helper object that manages the spatial/temporal bounds of the
78 task's quanta.
79 """
81 finders: dict[str, PrerequisiteFinder]
82 """Mapping of helper objects responsible for a single prerequisite input
83 connection.
85 Keys are connection names. Elements of this dictionary should be removed
86 by implementations of `QuantumGraphBuilder.process_subgraph` to take
87 responsibility for finding them away from the the `QuantumGraphBuilder`
88 base class.
89 """
91 def __init__(self, task_node: TaskNode, pipeline_graph: PipelineGraph):
92 self.bounds = PrerequisiteBounds(task_node)
93 self.finders = {
94 edge.connection_name: PrerequisiteFinder(edge, self.bounds, pipeline_graph)
95 for edge in task_node.prerequisite_inputs.values()
96 }
98 def update_bounds(self) -> None:
99 """Inspect the current state of `finders` and update `bounds` to
100 reflect the needs of only the finders that remain.
101 """
102 self.bounds.all_dataset_skypix.clear()
103 self.bounds.any_dataset_has_timespan = False
104 for finder in self.finders.values():
105 self.bounds.all_dataset_skypix.update(finder.dataset_skypix)
106 self.bounds.any_dataset_has_timespan = (
107 self.bounds.any_dataset_has_timespan or finder.dataset_has_timespan
108 )
111class PrerequisiteFinder:
112 """A QuantumGraph-generation helper class that manages the searches for a
113 prerequisite input connection.
115 Parameters
116 ----------
117 edge : `pipeline_graph.ReadEdge`
118 A `~pipeline_graph.PipelineGraph` edge that represents a single
119 prerequisite input connection.
120 bounds : `PrerequisiteBounds`
121 Another helper object that manages the spatial/temporal bounds of the
122 task's quanta, shared by all prerequisite inputs for that task.
123 pipeline_graph : `pipeline_graph.PipelineGraph`
124 Graph representation of the pipeline.
126 Notes
127 -----
128 `PrerequisiteFinder` instances are usually constructed by a
129 `PrerequisiteInfo` instance, which is in turn constructed by and attached
130 to the base `QuantumGraphBuilder` when a new builder is constructed. During
131 the `QuantumGraphBuilder.process_subgraph` hook implemented by a builder
132 subclass, prerequisite inputs may be found in other ways (e.g. via bulk
133 queries), as long as the results are consistent with the finder's
134 attributes, and this is indicated to the base `QuantumGraphBuilder` by
135 removing those finder instances after those prerequisites have been found
136 and added to a `QuantumGraphSkeleton`. Finder instances that remain in the
137 builder are used by calling `PrerequisiteFinder.find` on each quantum
138 later in `QuantumGraphBuilder.build`.
139 """
141 def __init__(
142 self,
143 edge: ReadEdge,
144 bounds: PrerequisiteBounds,
145 pipeline_graph: PipelineGraph,
146 ):
147 self.edge = edge
148 self._bounds = bounds
149 self.dataset_type_node = pipeline_graph.dataset_types[edge.parent_dataset_type_name]
150 self.lookup_function = self.task_node.get_lookup_function(edge.connection_name)
151 self.dataset_skypix = {}
152 self.dataset_other_spatial = {}
153 self.dataset_has_timespan = False
154 self.constraint_dimensions = self.task_node.dimensions
155 if self.lookup_function is None:
156 for family in self.dataset_type_node.dimensions.spatial - self.task_node.dimensions.spatial:
157 best_spatial_element = family.choose(
158 self.dataset_type_node.dimensions.elements.names,
159 self.dataset_type_node.dimensions.universe,
160 )
161 if isinstance(best_spatial_element, SkyPixDimension):
162 self.dataset_skypix[best_spatial_element.name] = best_spatial_element
163 else:
164 self.dataset_other_spatial[best_spatial_element.name] = cast(
165 DimensionElement, best_spatial_element
166 )
167 self.dataset_has_timespan = bool(
168 # If the task dimensions has a temporal family that isn't in
169 # the dataset type (i.e. "observation_timespans", like visit
170 # or exposure)...
171 self.task_node.dimensions.temporal
172 - self.dataset_type_node.dimensions.temporal
173 ) and (
174 # ...and the dataset type has a temporal family that isn't in
175 # the task dimensions, or is a calibration, the prerequisite
176 # search needs a temporal join. Note that the default
177 # dimension universe only has one temporal dimension family, so
178 # in practice this just means "calibration lookups when visit
179 # or exposure is in the task dimensions".
180 self.dataset_type_node.is_calibration
181 or bool(self.dataset_type_node.dimensions.temporal - self.task_node.dimensions.temporal)
182 )
183 new_constraint_dimensions = set()
184 universe = self.task_node.dimensions.universe
185 for dimension_name in self.task_node.dimensions.names:
186 if dimension_name in self.dataset_type_node.dimensions.names:
187 new_constraint_dimensions.add(dimension_name)
188 else:
189 dimension = universe[dimension_name]
190 if not (dimension.spatial or dimension.temporal):
191 new_constraint_dimensions.add(dimension_name)
192 self.constraint_dimensions = universe.conform(new_constraint_dimensions)
194 edge: ReadEdge
195 """The `~pipeline_graph.PipelineGraph` edge that represents the
196 prerequisite input connection.
197 """
199 dataset_type_node: DatasetTypeNode
200 """The `~pipeline_graph.PipelineGraph` node that represents the dataset
201 type of this connection.
203 This always uses the registry storage class and is never a component
204 dataset type.
205 """
207 lookup_function: Callable[
208 [DatasetType, Registry, DataCoordinate, Sequence[str]], Iterable[DatasetRef]
209 ] | None
210 """A task-provided callback for finding these datasets.
212 If this is not `None`, it must be used to ensure correct behavior.
213 """
215 dataset_skypix: dict[str, SkyPixDimension]
216 """Dimensions representing a pixelization of the sky used by the dataset
217 type for this connection that are also not part of the task's dimensions.
219 Keys are dimension names. It is at least extremely rare for this
220 dictionary to have more than one element.
221 """
223 dataset_other_spatial: dict[str, DimensionElement]
224 """Spatial dimensions other than sky pixelizations used by the dataset type
225 for this connection that are also not part of the task's dimensions.
226 """
228 dataset_has_timespan: bool
229 """Whether the dataset has a timespan that should be used in the lookup,
230 either because it is a calibration dataset or because it has temporal
231 dimensions that are not part of the tasks's dimensions.
232 """
234 @property
235 def task_node(self) -> TaskNode:
236 """The `~pipeline_graph.PipelineGraph` node that represents the task
237 for this connection.
238 """
239 return self._bounds.task_node
241 def find(
242 self,
243 butler: Butler,
244 input_collections: Sequence[str],
245 data_id: DataCoordinate,
246 skypix_bounds: Mapping[str, RangeSet],
247 timespan: Timespan | None,
248 ) -> list[DatasetRef]:
249 """Find prerequisite input datasets for a single quantum.
251 Parameters
252 ----------
253 butler : `lsst.daf.butler.Butler`
254 Butler client to use for queries.
255 input_collections : `~collections.abc.Sequence` [ `str` ]
256 Sequence of collections to search, in order.
257 data_id : `lsst.daf.butler.DataCoordinate`
258 Data ID for the quantum.
259 skypix_bounds : `Mapping` [ `str`, `lsst.sphgeom.RangeSet` ]
260 The spatial bounds of this quantum in various skypix dimensions.
261 Keys are skypix dimension names (a superset of those in
262 `dataset_skypix`) and values are sets of integer pixel ID ranges.
263 timespan : `lsst.daf.butler.Timespan` or `None`
264 The temporal bounds of this quantum. Guaranteed to not be `None`
265 if `dataset_has_timespan` is `True`.
267 Returns
268 -------
269 refs : `list` [ `lsst.daf.butler.DatasetRef` ]
270 Dataset references. These use
271 ``self.dataset_type_node.dataset_type``, which may differ from the
272 connection's dataset type in storage class or [lack of] component.
274 Raises
275 ------
276 NotImplementedError
277 Raised for certain relationships between task and dataset type
278 dimensions that are possible to define but not believed to be
279 useful in practice. These errors occur late rather than early in
280 order to allow a `QuantumGraphBuilder` subclass to handle them
281 first, in case an unusual task's needs must be met by a custom
282 builder class anyway.
283 """
284 if self.lookup_function:
285 # If there is a lookup function, just use it; nothing else matters.
286 return [
287 self.dataset_type_node.generalize_ref(ref)
288 for ref in self.lookup_function(
289 self.edge.adapt_dataset_type(self.dataset_type_node.dataset_type),
290 butler.registry,
291 data_id,
292 input_collections,
293 )
294 if ref is not None
295 ]
296 if self.dataset_type_node.dimensions <= self.constraint_dimensions:
297 # If this is a calibration dataset and the dataset doesn't have
298 # any dimensions that aren't constrained by the quantum data
299 # ID, we know there'll only be one result, and that means we
300 # can call Butler.find_dataset, which takes a timespan. Note
301 # that the AllDimensionsQuantumGraphBuilder subclass will
302 # intercept this case in order to optimize it when:
303 #
304 # - PipelineTaskConnections.getTemporalBoundsConnections is
305 # empty;
306 #
307 # - the quantum data IDs have temporal dimensions;
308 #
309 # and when that happens PrerequisiteFinder.find never gets
310 # called.
311 try:
312 ref = butler.find_dataset(
313 self.dataset_type_node.dataset_type,
314 data_id.subset(self.constraint_dimensions),
315 collections=input_collections,
316 timespan=timespan,
317 )
318 except MissingDatasetTypeError:
319 ref = None
320 return [ref] if ref is not None else []
321 elif self.dataset_has_timespan:
322 extra_dimensions = self.dataset_type_node.dimensions.names - self.constraint_dimensions.names
323 raise NotImplementedError(
324 f"No support for calibration lookup {self.task_node.label}.{self.edge.connection_name} "
325 f"with dimension(s) {extra_dimensions} not fully constrained by the task. "
326 "Please create a feature-request ticket and use a lookup function in the meantime."
327 )
328 if self.dataset_skypix:
329 if not self.dataset_has_timespan and not self.dataset_other_spatial:
330 # If the dataset has skypix dimensions but is not otherwise
331 # spatial or temporal (this describes reference catalogs and
332 # things like them), we can stuff the skypix IDs we want into
333 # the query via bind parameters and call queryDatasets. Once
334 # again AllDimensionsQuantumGraphBuilder will often intercept
335 # this case in order to optimize it, when:
336 #
337 # - PipelineTaskConnections.getSpatialBoundsConnections is
338 # empty;
339 #
340 # - the quantum data IDs have spatial dimensions;
341 #
342 # and when that happens PrerequisiteFinder.find never gets
343 # called.
344 where_terms: list[str] = []
345 bind: dict[str, list[int]] = {}
346 for name in self.dataset_skypix:
347 where_terms.append(f"{name} IN ({name}_pixels)")
348 pixels: list[int] = []
349 for begin, end in skypix_bounds[name]:
350 pixels.extend(range(begin, end))
351 bind[f"{name}_pixels"] = pixels
352 try:
353 return list(
354 butler.registry.queryDatasets(
355 self.dataset_type_node.dataset_type,
356 collections=input_collections,
357 dataId=data_id.subset(self.constraint_dimensions),
358 where=" AND ".join(where_terms),
359 bind=bind,
360 findFirst=True,
361 ).expanded()
362 )
363 except MissingDatasetTypeError:
364 return []
365 else:
366 raise NotImplementedError(
367 f"No support for skypix lookup {self.task_node.label}.{self.edge.connection_name} "
368 "that requires additional spatial and/or temporal constraints. "
369 "Please create a feature-request ticket and use a lookup function in the meantime."
370 )
371 if self._bounds.spatial_connections or self._bounds.temporal_connections:
372 raise NotImplementedError(
373 f"No support for prerequisite lookup {self.task_node.label}.{self.edge.connection_name} "
374 "that requires other connections to determine spatial or temporal bounds but does not "
375 "fit into one of our standard cases. "
376 "Please create a feature-request ticket and use a lookup function in the meantime."
377 )
378 # If the spatial/temporal bounds are not customized, and the dataset
379 # doesn't have any skypix dimensions, a vanilla queryDatasets call
380 # should work. This case should always be optimized by
381 # AllDimensionsQuantumGraphBuilder as well. Note that we use the
382 # original quantum data ID here, not those with constraint_dimensions
383 # that strips out the spatial/temporal stuff, because here we want the
384 # butler query system to handle the spatial/temporal stuff like it
385 # normally would.
386 try:
387 return list(
388 butler.registry.queryDatasets(
389 self.dataset_type_node.dataset_type,
390 collections=input_collections,
391 dataId=data_id,
392 findFirst=True,
393 ).expanded()
394 )
395 except MissingDatasetTypeError:
396 return []
399@dataclasses.dataclass
400class PrerequisiteBounds:
401 """A QuantumGraph-generation helper class that manages the spatial and
402 temporal bounds of a tasks' quanta, for the purpose of finding
403 prerequisite inputs.
404 """
406 task_node: TaskNode
407 """The `~pipeline_graph.PipelineGraph` node that represents the task."""
409 spatial_connections: frozenset[str] = dataclasses.field(init=False)
410 """Regular input or output connections whose (assumed spatial) data IDs
411 should be used to define the spatial bounds of this task's quanta.
413 See Also
414 --------
415 PipelineTaskConnections.getSpatialBoundsConnections
416 """
418 temporal_connections: frozenset[str] = dataclasses.field(init=False)
419 """Regular input or output connections whose (assumed temporal) data IDs
420 should be used to define the temporal bounds of this task's quanta.
422 See Also
423 --------
424 PipelineTaskConnections.getTemporalBoundsConnections
425 """
427 all_dataset_skypix: dict[str, SkyPixDimension] = dataclasses.field(default_factory=dict)
428 """The union of all `PrerequisiteFinder.dataset_skypix` attributes for all
429 (remaining) prerequisite finders for this task.
430 """
432 any_dataset_has_timespan: bool = dataclasses.field(default=False)
433 """Whether any `PrerequisiteFinder.dataset_has_timespan` attribute is true
434 for any (remaining) prerequisite finder for this task.
435 """
437 def __post_init__(self) -> None:
438 self.spatial_connections = frozenset(self.task_node.get_spatial_bounds_connections())
439 self.temporal_connections = frozenset(self.task_node.get_temporal_bounds_connections())
441 def make_skypix_bounds_builder(self, quantum_data_id: DataCoordinate) -> SkyPixBoundsBuilder:
442 """Return an object that accumulates the appropriate spatial bounds for
443 a quantum.
445 Parameters
446 ----------
447 quantum_data_id : `lsst.daf.butler.DataCoordinate`
448 Data ID for this quantum.
450 Returns
451 -------
452 builder : `SkyPixBoundsBuilder`
453 Object that accumulates the appropriate spatial bounds for a
454 quantum. If the spatial bounds are not needed, this object will do
455 nothing.
456 """
457 if not self.all_dataset_skypix:
458 return _TrivialSkyPixBoundsBuilder()
459 if self.spatial_connections:
460 return _ConnectionSkyPixBoundsBuilder(
461 self.task_node, self.spatial_connections, self.all_dataset_skypix.values(), quantum_data_id
462 )
463 if self.task_node.dimensions.spatial:
464 return _QuantumOnlySkyPixBoundsBuilder(self.all_dataset_skypix.values(), quantum_data_id)
465 else:
466 return _UnboundedSkyPixBoundsBuilder(self.all_dataset_skypix.values())
468 def make_timespan_builder(self, quantum_data_id: DataCoordinate) -> TimespanBuilder:
469 """Return an object that accumulates the appropriate timespan for
470 a quantum.
472 Parameters
473 ----------
474 quantum_data_id : `lsst.daf.butler.DataCoordinate`
475 Data ID for this quantum.
477 Returns
478 -------
479 builder : `TimespanBuilder`
480 Object that accumulates the appropriate timespan bounds for a
481 quantum. If a timespan is not needed, this object will do nothing.
482 """
483 if not self.any_dataset_has_timespan:
484 return _TrivialTimespanBuilder()
485 if self.temporal_connections:
486 return _ConnectionTimespanBuilder(self.task_node, self.temporal_connections, quantum_data_id)
487 if self.task_node.dimensions.temporal:
488 return _QuantumOnlyTimespanBuilder(quantum_data_id)
489 else:
490 return _UnboundedTimespanBuilder()
493class SkyPixBoundsBuilder(ABC):
494 """A base class for objects that accumulate the appropriate spatial bounds
495 for a quantum.
496 """
498 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
499 """Handle the skeleton graph node for a regular input/output connection
500 for this quantum, including its data ID in the bounds if appropriate.
502 Parameters
503 ----------
504 parent_dataset_type_name : `str`
505 Name of the dataset type. Never a component dataset type name.
506 data_id : `lsst.daf.butler.DataCoordinate`
507 Data ID for the dataset.
508 """
509 pass
511 @abstractmethod
512 def finish(self) -> dict[str, RangeSet]:
513 """Finish building the spatial bounds and return them.
515 Returns
516 -------
517 bounds : `dict` [ `str`, `lsst.sphgeom.RangeSet` ]
518 The spatial bounds of this quantum in various skypix dimensions.
519 Keys are skypix dimension names and values are sets of integer
520 pixel ID ranges.
521 """
522 raise NotImplementedError()
525class TimespanBuilder(ABC):
526 """A base class for objects that accumulate the appropriate timespan
527 for a quantum.
528 """
530 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
531 """Handle the skeleton graph node for a regular input/output connection
532 for this quantum, including its data ID in the bounds if appropriate.
534 Parameters
535 ----------
536 parent_dataset_type_name : `str`
537 Name of the dataset type. Never a component dataset type name.
538 data_id : `lsst.daf.butler.DataCoordinate`
539 Data ID for the dataset.
540 """
541 pass
543 @abstractmethod
544 def finish(self) -> Timespan | None:
545 """Finish building the timespan and return it.
547 Returns
548 -------
549 timespan : `lsst.daf.butler.Timespan` or `None`
550 The timespan of this quantum, or `None` if it is known to not be
551 needed.
552 """
553 raise NotImplementedError()
556class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder):
557 """Implementation of `SkyPixBoundsBuilder` for when no skypix bounds are
558 needed.
559 """
561 def finish(self) -> dict[str, RangeSet]:
562 return {}
565class _TrivialTimespanBuilder(TimespanBuilder):
566 """Implementation of `TimespanBuilder` for when no timespan is needed."""
568 def finish(self) -> None:
569 return None
572class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder):
573 """Implementation of `SkyPixBoundsBuilder` for when the quantum data IDs
574 provide the only relevant spatial regions.
575 """
577 def __init__(self, dimensions: Iterable[SkyPixDimension], quantum_data_id: DataCoordinate) -> None:
578 self._region = quantum_data_id.region
579 self._dimensions = dimensions
581 def finish(self) -> dict[str, RangeSet]:
582 return {
583 dimension.name: dimension.pixelization.envelope(self._region) for dimension in self._dimensions
584 }
587class _QuantumOnlyTimespanBuilder(TimespanBuilder):
588 """Implementation of `TimespanBuilder` for when the quantum data IDs
589 provide the only relevant timespans.
590 """
592 def __init__(self, quantum_data_id: DataCoordinate) -> None:
593 self._timespan = cast(Timespan, quantum_data_id.timespan)
595 def finish(self) -> Timespan:
596 return self._timespan
599class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder):
600 """Implementation of `SkyPixBoundsBuilder` for when the bounds cover the
601 full sky.
602 """
604 def __init__(self, dimensions: Iterable[SkyPixDimension]):
605 self._dimensions = dimensions
607 def finish(self) -> dict[str, RangeSet]:
608 return {dimension.name: dimension.pixelization.universe() for dimension in self._dimensions}
611class _UnboundedTimespanBuilder(TimespanBuilder):
612 """Implementation of `TimespanBuilder` for when the timespan covers all
613 time.
614 """
616 def finish(self) -> Timespan:
617 return Timespan(None, None)
620class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder):
621 """Implementation of `SkyPixBoundsBuilder` for when other input or output
622 connections contribute to the spatial bounds.
623 """
625 def __init__(
626 self,
627 task_node: TaskNode,
628 bounds_connections: frozenset[str],
629 dimensions: Iterable[SkyPixDimension],
630 quantum_data_id: DataCoordinate,
631 ) -> None:
632 self._dimensions = dimensions
633 self._regions: list[Region] = []
634 if task_node.dimensions.spatial:
635 self._regions.append(quantum_data_id.region)
636 self._dataset_type_names: set[str] = set()
637 for connection_name in bounds_connections:
638 if edge := task_node.inputs.get(connection_name):
639 self._dataset_type_names.add(edge.parent_dataset_type_name)
640 else:
641 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name)
642 # Note that we end up raising if the input is a prerequisite (and
643 # hence not in task_node.inputs or task_node.outputs); this
644 # justifies the cast in `handle_dataset`.
646 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
647 if parent_dataset_type_name in self._dataset_type_names:
648 self._regions.append(data_id.region)
650 def finish(self) -> dict[str, RangeSet]:
651 result = {}
652 for dimension in self._dimensions:
653 bounds = RangeSet()
654 for region in self._regions:
655 bounds |= dimension.pixelization.envelope(region)
656 result[dimension.name] = bounds
657 return result
660class _ConnectionTimespanBuilder(TimespanBuilder):
661 """Implementation of `TimespanBuilder` for when other input or output
662 connections contribute to the timespan.
663 """
665 def __init__(
666 self,
667 task_node: TaskNode,
668 bounds_connections: frozenset[str],
669 quantum_data_id: DataCoordinate,
670 ) -> None:
671 timespan = (
672 cast(Timespan, quantum_data_id.timespan)
673 if task_node.dimensions.temporal
674 else Timespan.makeEmpty()
675 )
676 self._begin_nsec = timespan._nsec[0]
677 self._end_nsec = timespan._nsec[1]
678 self._dataset_type_names = set()
679 for connection_name in bounds_connections:
680 if edge := task_node.inputs.get(connection_name):
681 self._dataset_type_names.add(edge.parent_dataset_type_name)
682 else:
683 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name)
684 # Note that we end up raising if the input is a prerequisite (and
685 # hence not in task_node.inputs or task_node.outputs); this
686 # justifies the cast in `handle_dataset`.
688 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
689 if parent_dataset_type_name in self._dataset_type_names:
690 nsec = cast(Timespan, data_id.timespan)._nsec
691 self._begin_nsec = min(self._begin_nsec, nsec[0])
692 self._end_nsec = max(self._end_nsec, nsec[1])
694 def finish(self) -> Timespan:
695 return Timespan(None, None, _nsec=(self._begin_nsec, self._end_nsec))