Coverage for python / lsst / pipe / base / prerequisite_helpers.py: 32%
201 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-17 16:30 -0800
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-17 16:30 -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/>.
28"""Helper classes for finding prerequisite input datasets during
29QuantumGraph generation.
30"""
32from __future__ import annotations
34__all__ = (
35 "PrerequisiteBounds",
36 "PrerequisiteFinder",
37 "PrerequisiteInfo",
38 "SkyPixBoundsBuilder",
39 "TimespanBuilder",
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(self.dataset_type_node.dimensions)
158 if isinstance(best_spatial_element, SkyPixDimension):
159 self.dataset_skypix[best_spatial_element.name] = best_spatial_element
160 else:
161 self.dataset_other_spatial[best_spatial_element.name] = cast(
162 DimensionElement, best_spatial_element
163 )
164 self.dataset_has_timespan = bool(
165 # If the task dimensions has a temporal family that isn't in
166 # the dataset type (i.e. "observation_timespans", like visit
167 # or exposure)...
168 self.task_node.dimensions.temporal - self.dataset_type_node.dimensions.temporal
169 ) and (
170 # ...and the dataset type has a temporal family that isn't in
171 # the task dimensions, or is a calibration, the prerequisite
172 # search needs a temporal join. Note that the default
173 # dimension universe only has one temporal dimension family, so
174 # in practice this just means "calibration lookups when visit
175 # or exposure is in the task dimensions".
176 self.dataset_type_node.is_calibration
177 or bool(self.dataset_type_node.dimensions.temporal - self.task_node.dimensions.temporal)
178 )
179 new_constraint_dimensions = set()
180 universe = self.task_node.dimensions.universe
181 for dimension_name in self.task_node.dimensions.names:
182 if dimension_name in self.dataset_type_node.dimensions.names:
183 new_constraint_dimensions.add(dimension_name)
184 else:
185 dimension = universe[dimension_name]
186 if not (dimension.spatial or dimension.temporal):
187 new_constraint_dimensions.add(dimension_name)
188 self.constraint_dimensions = universe.conform(new_constraint_dimensions)
190 edge: ReadEdge
191 """The `~pipeline_graph.PipelineGraph` edge that represents the
192 prerequisite input connection.
193 """
195 dataset_type_node: DatasetTypeNode
196 """The `~pipeline_graph.PipelineGraph` node that represents the dataset
197 type of this connection.
199 This always uses the registry storage class and is never a component
200 dataset type.
201 """
203 lookup_function: (
204 Callable[[DatasetType, Registry, DataCoordinate, Sequence[str]], Iterable[DatasetRef]] | None
205 )
206 """A task-provided callback for finding these datasets.
208 If this is not `None`, it must be used to ensure correct behavior.
209 """
211 dataset_skypix: dict[str, SkyPixDimension]
212 """Dimensions representing a pixelization of the sky used by the dataset
213 type for this connection that are also not part of the task's dimensions.
215 Keys are dimension names. It is at least extremely rare for this
216 dictionary to have more than one element.
217 """
219 dataset_other_spatial: dict[str, DimensionElement]
220 """Spatial dimensions other than sky pixelizations used by the dataset type
221 for this connection that are also not part of the task's dimensions.
222 """
224 dataset_has_timespan: bool
225 """Whether the dataset has a timespan that should be used in the lookup,
226 either because it is a calibration dataset or because it has temporal
227 dimensions that are not part of the tasks's dimensions.
228 """
230 @property
231 def task_node(self) -> TaskNode:
232 """The `~pipeline_graph.PipelineGraph` node that represents the task
233 for this connection.
234 """
235 return self._bounds.task_node
237 def find(
238 self,
239 butler: Butler,
240 input_collections: Sequence[str],
241 data_id: DataCoordinate,
242 skypix_bounds: Mapping[str, RangeSet],
243 timespan: Timespan | None,
244 ) -> list[DatasetRef]:
245 """Find prerequisite input datasets for a single quantum.
247 Parameters
248 ----------
249 butler : `lsst.daf.butler.Butler`
250 Butler client to use for queries.
251 input_collections : `~collections.abc.Sequence` [ `str` ]
252 Sequence of collections to search, in order.
253 data_id : `lsst.daf.butler.DataCoordinate`
254 Data ID for the quantum.
255 skypix_bounds : `Mapping` [ `str`, `lsst.sphgeom.RangeSet` ]
256 The spatial bounds of this quantum in various skypix dimensions.
257 Keys are skypix dimension names (a superset of those in
258 `dataset_skypix`) and values are sets of integer pixel ID ranges.
259 timespan : `lsst.daf.butler.Timespan` or `None`
260 The temporal bounds of this quantum. Guaranteed to not be `None`
261 if `dataset_has_timespan` is `True`.
263 Returns
264 -------
265 refs : `list` [ `lsst.daf.butler.DatasetRef` ]
266 Dataset references. These use
267 ``self.dataset_type_node.dataset_type``, which may differ from the
268 connection's dataset type in storage class or [lack of] component.
270 Raises
271 ------
272 NotImplementedError
273 Raised for certain relationships between task and dataset type
274 dimensions that are possible to define but not believed to be
275 useful in practice. These errors occur late rather than early in
276 order to allow a `QuantumGraphBuilder` subclass to handle them
277 first, in case an unusual task's needs must be met by a custom
278 builder class anyway.
279 """
280 if self.lookup_function:
281 # If there is a lookup function, just use it; nothing else matters.
282 return [
283 self.dataset_type_node.generalize_ref(ref)
284 for ref in self.lookup_function(
285 self.edge.adapt_dataset_type(self.dataset_type_node.dataset_type),
286 butler.registry,
287 data_id,
288 input_collections,
289 )
290 if ref is not None
291 ]
292 if self.dataset_type_node.dimensions <= self.constraint_dimensions:
293 # If this is a calibration dataset and the dataset doesn't have
294 # any dimensions that aren't constrained by the quantum data
295 # ID, we know there'll only be one result, and that means we
296 # can call Butler.find_dataset, which takes a timespan. Note
297 # that the AllDimensionsQuantumGraphBuilder subclass will
298 # intercept this case in order to optimize it when:
299 #
300 # - PipelineTaskConnections.getTemporalBoundsConnections is
301 # empty;
302 #
303 # - the quantum data IDs have temporal dimensions;
304 #
305 # and when that happens PrerequisiteFinder.find never gets
306 # called.
307 try:
308 ref = butler.find_dataset(
309 self.dataset_type_node.dataset_type,
310 data_id.subset(self.constraint_dimensions),
311 collections=input_collections,
312 timespan=timespan,
313 )
314 except MissingDatasetTypeError:
315 ref = None
316 return [ref] if ref is not None else []
317 elif self.dataset_has_timespan:
318 extra_dimensions = self.dataset_type_node.dimensions.names - self.constraint_dimensions.names
319 raise NotImplementedError(
320 f"No support for calibration lookup {self.task_node.label}.{self.edge.connection_name} "
321 f"with dimension(s) {extra_dimensions} not fully constrained by the task. "
322 "Please create a feature-request ticket and use a lookup function in the meantime."
323 )
324 if self.dataset_skypix:
325 if not self.dataset_has_timespan and not self.dataset_other_spatial:
326 # If the dataset has skypix dimensions but is not otherwise
327 # spatial or temporal (this describes reference catalogs and
328 # things like them), we can stuff the skypix IDs we want into
329 # the query via bind parameters and call queryDatasets. Once
330 # again AllDimensionsQuantumGraphBuilder will often intercept
331 # this case in order to optimize it, when:
332 #
333 # - PipelineTaskConnections.getSpatialBoundsConnections is
334 # empty;
335 #
336 # - the quantum data IDs have spatial dimensions;
337 #
338 # and when that happens PrerequisiteFinder.find never gets
339 # called.
340 where_terms: list[str] = []
341 bind: dict[str, list[int]] = {}
342 for name in self.dataset_skypix:
343 where_terms.append(f"{name} IN (:{name}_pixels)")
344 pixels: list[int] = []
345 for begin, end in skypix_bounds[name]:
346 pixels.extend(range(begin, end))
347 bind[f"{name}_pixels"] = pixels
348 try:
349 return butler.query_datasets(
350 self.dataset_type_node.dataset_type,
351 collections=input_collections,
352 data_id=data_id.subset(self.constraint_dimensions),
353 where=" AND ".join(where_terms),
354 bind=bind,
355 with_dimension_records=True,
356 limit=None,
357 explain=False,
358 )
359 except MissingDatasetTypeError:
360 return []
361 else:
362 raise NotImplementedError(
363 f"No support for skypix lookup {self.task_node.label}.{self.edge.connection_name} "
364 "that requires additional spatial and/or temporal constraints. "
365 "Please create a feature-request ticket and use a lookup function in the meantime."
366 )
367 if self._bounds.spatial_connections or self._bounds.temporal_connections:
368 raise NotImplementedError(
369 f"No support for prerequisite lookup {self.task_node.label}.{self.edge.connection_name} "
370 "that requires other connections to determine spatial or temporal bounds but does not "
371 "fit into one of our standard cases. "
372 "Please create a feature-request ticket and use a lookup function in the meantime."
373 )
374 # If the spatial/temporal bounds are not customized, and the dataset
375 # doesn't have any skypix dimensions, a vanilla query_datasets call
376 # should work. This case should always be optimized by
377 # AllDimensionsQuantumGraphBuilder as well. Note that we use the
378 # original quantum data ID here, not those with constraint_dimensions
379 # that strips out the spatial/temporal stuff, because here we want the
380 # butler query system to handle the spatial/temporal stuff like it
381 # normally would.
382 try:
383 return butler.query_datasets(
384 self.dataset_type_node.dataset_type,
385 collections=input_collections,
386 data_id=data_id,
387 with_dimension_records=True,
388 limit=None,
389 explain=False,
390 )
391 except MissingDatasetTypeError:
392 return []
395@dataclasses.dataclass
396class PrerequisiteBounds:
397 """A QuantumGraph-generation helper class that manages the spatial and
398 temporal bounds of a tasks' quanta, for the purpose of finding
399 prerequisite inputs.
400 """
402 task_node: TaskNode
403 """The `~pipeline_graph.PipelineGraph` node that represents the task."""
405 spatial_connections: frozenset[str] = dataclasses.field(init=False)
406 """Regular input or output connections whose (assumed spatial) data IDs
407 should be used to define the spatial bounds of this task's quanta.
409 See Also
410 --------
411 PipelineTaskConnections.getSpatialBoundsConnections
412 """
414 temporal_connections: frozenset[str] = dataclasses.field(init=False)
415 """Regular input or output connections whose (assumed temporal) data IDs
416 should be used to define the temporal bounds of this task's quanta.
418 See Also
419 --------
420 PipelineTaskConnections.getTemporalBoundsConnections
421 """
423 all_dataset_skypix: dict[str, SkyPixDimension] = dataclasses.field(default_factory=dict)
424 """The union of all `PrerequisiteFinder.dataset_skypix` attributes for all
425 (remaining) prerequisite finders for this task.
426 """
428 any_dataset_has_timespan: bool = dataclasses.field(default=False)
429 """Whether any `PrerequisiteFinder.dataset_has_timespan` attribute is true
430 for any (remaining) prerequisite finder for this task.
431 """
433 def __post_init__(self) -> None:
434 self.spatial_connections = frozenset(self.task_node.get_spatial_bounds_connections())
435 self.temporal_connections = frozenset(self.task_node.get_temporal_bounds_connections())
437 def make_skypix_bounds_builder(self, quantum_data_id: DataCoordinate) -> SkyPixBoundsBuilder:
438 """Return an object that accumulates the appropriate spatial bounds for
439 a quantum.
441 Parameters
442 ----------
443 quantum_data_id : `lsst.daf.butler.DataCoordinate`
444 Data ID for this quantum.
446 Returns
447 -------
448 builder : `SkyPixBoundsBuilder`
449 Object that accumulates the appropriate spatial bounds for a
450 quantum. If the spatial bounds are not needed, this object will do
451 nothing.
452 """
453 if not self.all_dataset_skypix:
454 return _TrivialSkyPixBoundsBuilder()
455 if self.spatial_connections:
456 return _ConnectionSkyPixBoundsBuilder(
457 self.task_node, self.spatial_connections, self.all_dataset_skypix.values(), quantum_data_id
458 )
459 if self.task_node.dimensions.spatial:
460 return _QuantumOnlySkyPixBoundsBuilder(self.all_dataset_skypix.values(), quantum_data_id)
461 else:
462 return _UnboundedSkyPixBoundsBuilder(self.all_dataset_skypix.values())
464 def make_timespan_builder(self, quantum_data_id: DataCoordinate) -> TimespanBuilder:
465 """Return an object that accumulates the appropriate timespan for
466 a quantum.
468 Parameters
469 ----------
470 quantum_data_id : `lsst.daf.butler.DataCoordinate`
471 Data ID for this quantum.
473 Returns
474 -------
475 builder : `TimespanBuilder`
476 Object that accumulates the appropriate timespan bounds for a
477 quantum. If a timespan is not needed, this object will do nothing.
478 """
479 if not self.any_dataset_has_timespan:
480 return _TrivialTimespanBuilder()
481 if self.temporal_connections:
482 return _ConnectionTimespanBuilder(self.task_node, self.temporal_connections, quantum_data_id)
483 if self.task_node.dimensions.temporal:
484 return _QuantumOnlyTimespanBuilder(quantum_data_id)
485 else:
486 return _UnboundedTimespanBuilder()
489class SkyPixBoundsBuilder(ABC):
490 """A base class for objects that accumulate the appropriate spatial bounds
491 for a quantum.
492 """
494 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
495 """Handle the skeleton graph node for a regular input/output connection
496 for this quantum, including its data ID in the bounds if appropriate.
498 Parameters
499 ----------
500 parent_dataset_type_name : `str`
501 Name of the dataset type. Never a component dataset type name.
502 data_id : `lsst.daf.butler.DataCoordinate`
503 Data ID for the dataset.
504 """
505 pass
507 @abstractmethod
508 def finish(self) -> dict[str, RangeSet]:
509 """Finish building the spatial bounds and return them.
511 Returns
512 -------
513 bounds : `dict` [ `str`, `lsst.sphgeom.RangeSet` ]
514 The spatial bounds of this quantum in various skypix dimensions.
515 Keys are skypix dimension names and values are sets of integer
516 pixel ID ranges.
517 """
518 raise NotImplementedError()
521class TimespanBuilder(ABC):
522 """A base class for objects that accumulate the appropriate timespan
523 for a quantum.
524 """
526 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
527 """Handle the skeleton graph node for a regular input/output connection
528 for this quantum, including its data ID in the bounds if appropriate.
530 Parameters
531 ----------
532 parent_dataset_type_name : `str`
533 Name of the dataset type. Never a component dataset type name.
534 data_id : `lsst.daf.butler.DataCoordinate`
535 Data ID for the dataset.
536 """
537 pass
539 @abstractmethod
540 def finish(self) -> Timespan | None:
541 """Finish building the timespan and return it.
543 Returns
544 -------
545 timespan : `lsst.daf.butler.Timespan` or `None`
546 The timespan of this quantum, or `None` if it is known to not be
547 needed.
548 """
549 raise NotImplementedError()
552class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder):
553 """Implementation of `SkyPixBoundsBuilder` for when no skypix bounds are
554 needed.
555 """
557 def finish(self) -> dict[str, RangeSet]:
558 return {}
561class _TrivialTimespanBuilder(TimespanBuilder):
562 """Implementation of `TimespanBuilder` for when no timespan is needed."""
564 def finish(self) -> None:
565 return None
568class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder):
569 """Implementation of `SkyPixBoundsBuilder` for when the quantum data IDs
570 provide the only relevant spatial regions.
571 """
573 def __init__(self, dimensions: Iterable[SkyPixDimension], quantum_data_id: DataCoordinate) -> None:
574 self._region = quantum_data_id.region
575 self._dimensions = dimensions
577 def finish(self) -> dict[str, RangeSet]:
578 return {
579 dimension.name: dimension.pixelization.envelope(self._region) for dimension in self._dimensions
580 }
583class _QuantumOnlyTimespanBuilder(TimespanBuilder):
584 """Implementation of `TimespanBuilder` for when the quantum data IDs
585 provide the only relevant timespans.
586 """
588 def __init__(self, quantum_data_id: DataCoordinate) -> None:
589 self._timespan = cast(Timespan, quantum_data_id.timespan)
591 def finish(self) -> Timespan:
592 return self._timespan
595class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder):
596 """Implementation of `SkyPixBoundsBuilder` for when the bounds cover the
597 full sky.
598 """
600 def __init__(self, dimensions: Iterable[SkyPixDimension]):
601 self._dimensions = dimensions
603 def finish(self) -> dict[str, RangeSet]:
604 return {dimension.name: dimension.pixelization.universe() for dimension in self._dimensions}
607class _UnboundedTimespanBuilder(TimespanBuilder):
608 """Implementation of `TimespanBuilder` for when the timespan covers all
609 time.
610 """
612 def finish(self) -> Timespan:
613 return Timespan(None, None)
616class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder):
617 """Implementation of `SkyPixBoundsBuilder` for when other input or output
618 connections contribute to the spatial bounds.
619 """
621 def __init__(
622 self,
623 task_node: TaskNode,
624 bounds_connections: frozenset[str],
625 dimensions: Iterable[SkyPixDimension],
626 quantum_data_id: DataCoordinate,
627 ) -> None:
628 self._dimensions = dimensions
629 self._regions: list[Region] = []
630 if task_node.dimensions.spatial:
631 self._regions.append(quantum_data_id.region)
632 self._dataset_type_names: set[str] = set()
633 for connection_name in bounds_connections:
634 if edge := task_node.inputs.get(connection_name):
635 self._dataset_type_names.add(edge.parent_dataset_type_name)
636 else:
637 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name)
638 # Note that we end up raising if the input is a prerequisite (and
639 # hence not in task_node.inputs or task_node.outputs); this
640 # justifies the cast in `handle_dataset`.
642 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
643 if parent_dataset_type_name in self._dataset_type_names:
644 self._regions.append(data_id.region)
646 def finish(self) -> dict[str, RangeSet]:
647 result = {}
648 for dimension in self._dimensions:
649 bounds = RangeSet()
650 for region in self._regions:
651 bounds |= dimension.pixelization.envelope(region)
652 result[dimension.name] = bounds
653 return result
656class _ConnectionTimespanBuilder(TimespanBuilder):
657 """Implementation of `TimespanBuilder` for when other input or output
658 connections contribute to the timespan.
659 """
661 def __init__(
662 self,
663 task_node: TaskNode,
664 bounds_connections: frozenset[str],
665 quantum_data_id: DataCoordinate,
666 ) -> None:
667 timespan = (
668 cast(Timespan, quantum_data_id.timespan)
669 if task_node.dimensions.temporal
670 else Timespan.makeEmpty()
671 )
672 self._begin_nsec = timespan.nsec[0]
673 self._end_nsec = timespan.nsec[1]
674 self._dataset_type_names = set()
675 for connection_name in bounds_connections:
676 if edge := task_node.inputs.get(connection_name):
677 self._dataset_type_names.add(edge.parent_dataset_type_name)
678 else:
679 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name)
680 # Note that we end up raising if the input is a prerequisite (and
681 # hence not in task_node.inputs or task_node.outputs); this
682 # justifies the cast in `handle_dataset`.
684 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None:
685 if parent_dataset_type_name in self._dataset_type_names:
686 nsec = cast(Timespan, data_id.timespan).nsec
687 self._begin_nsec = min(self._begin_nsec, nsec[0])
688 self._end_nsec = max(self._end_nsec, nsec[1])
690 def finish(self) -> Timespan:
691 return Timespan(None, None, _nsec=(self._begin_nsec, self._end_nsec))