Coverage for python/lsst/pipe/base/prerequisite_helpers.py: 35%

194 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-17 10:52 +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/>. 

27 

28"""Helper classes for finding prerequisite input datasets during 

29QuantumGraph generation. 

30""" 

31 

32from __future__ import annotations 

33 

34__all__ = ( 

35 "SkyPixBoundsBuilder", 

36 "PrerequisiteFinder", 

37 "PrerequisiteBounds", 

38 "TimespanBuilder", 

39 "PrerequisiteInfo", 

40) 

41 

42import dataclasses 

43from abc import ABC, abstractmethod 

44from collections.abc import Callable, Iterable, Mapping, Sequence 

45from typing import cast 

46 

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 

59 

60from .pipeline_graph import DatasetTypeNode, PipelineGraph, ReadEdge, TaskNode 

61 

62 

63@dataclasses.dataclass 

64class PrerequisiteInfo: 

65 """A QuantumGraph-generation helper class that manages the searches for all 

66 prerequisite input connections for a task. 

67 """ 

68 

69 bounds: PrerequisiteBounds 

70 """Another helper object that manages the spatial/temporal bounds of the 

71 task's quanta. 

72 """ 

73 

74 finders: dict[str, PrerequisiteFinder] 

75 """Mapping of helper objects responsible for a single prerequisite input 

76 connection. 

77 

78 Keys are connection names. Elements of this dictionary should be removed 

79 by implementations of `QuantumGraphBuilder.process_subgraph` to take 

80 responsibility for finding them away from the the `QuantumGraphBuilder` 

81 base class. 

82 """ 

83 

84 def __init__(self, task_node: TaskNode, pipeline_graph: PipelineGraph): 

85 self.bounds = PrerequisiteBounds(task_node) 

86 self.finders = { 

87 edge.connection_name: PrerequisiteFinder(edge, self.bounds, pipeline_graph) 

88 for edge in task_node.prerequisite_inputs.values() 

89 } 

90 

91 def update_bounds(self) -> None: 

92 """Inspect the current state of `finders` and update `bounds` to 

93 reflect the needs of only the finders that remain. 

94 """ 

95 self.bounds.all_dataset_skypix.clear() 

96 self.bounds.any_dataset_has_timespan = False 

97 for finder in self.finders.values(): 

98 self.bounds.all_dataset_skypix.update(finder.dataset_skypix) 

99 self.bounds.any_dataset_has_timespan = ( 

100 self.bounds.any_dataset_has_timespan or finder.dataset_has_timespan 

101 ) 

102 

103 

104class PrerequisiteFinder: 

105 """A QuantumGraph-generation helper class that manages the searches for a 

106 prerequisite input connection. 

107 

108 Parameters 

109 ---------- 

110 edge : `pipeline_graph.ReadEdge` 

111 A `~pipeline_graph.PipelineGraph` edge that represents a single 

112 prerequisite input connection. 

113 bounds : `PrerequisiteBounds` 

114 Another helper object that manages the spatial/temporal bounds of the 

115 task's quanta, shared by all prerequisite inputs for that task. 

116 pipeline_graph `pipeline_graph.PipelineGraph` 

117 Graph representation of the pipeline. 

118 

119 Notes 

120 ----- 

121 `PrerequisiteFinder` instances are usually constructed by a 

122 `PrerequisiteInfo` instance, which is in turn constructed by and attached 

123 to the base `QuantumGraphBuilder` when a new builder is constructed. During 

124 the `QuantumGraphBuilder.process_subgraph` hook implemented by a builder 

125 subclass, prerequisite inputs may be found in other ways (e.g. via bulk 

126 queries), as long as the results are consistent with the finder's 

127 attributes, and this is indicated to the base `QuantumGraphBuilder` by 

128 removing those finder instances after those prerequisites have been found 

129 and added to a `QuantumGraphSkeleton`. Finder instances that remain in the 

130 builder are used by calling `PrerequisiteFinder.find` on each quantum 

131 later in `QuantumGraphBuilder.build`. 

132 """ 

133 

134 def __init__( 

135 self, 

136 edge: ReadEdge, 

137 bounds: PrerequisiteBounds, 

138 pipeline_graph: PipelineGraph, 

139 ): 

140 self.edge = edge 

141 self._bounds = bounds 

142 self.dataset_type_node = pipeline_graph.dataset_types[edge.parent_dataset_type_name] 

143 self.lookup_function = self.task_node.get_lookup_function(edge.connection_name) 

144 self.dataset_skypix = {} 

145 self.dataset_other_spatial = {} 

146 self.dataset_has_timespan = False 

147 self.constraint_dimensions = self.task_node.dimensions 

148 if self.lookup_function is None: 

149 for family in self.dataset_type_node.dimensions.spatial - self.task_node.dimensions.spatial: 

150 best_spatial_element = family.choose(self.dataset_type_node.dimensions.elements) 

151 if isinstance(best_spatial_element, SkyPixDimension): 

152 self.dataset_skypix[best_spatial_element.name] = best_spatial_element 

153 else: 

154 self.dataset_other_spatial[best_spatial_element.name] = cast( 

155 DimensionElement, best_spatial_element 

156 ) 

157 self.dataset_has_timespan = self.dataset_type_node.is_calibration or bool( 

158 self.dataset_type_node.dimensions.temporal - self.task_node.dimensions.temporal 

159 ) 

160 self.constraint_dimensions = self.constraint_dimensions.universe.extract( 

161 { 

162 d.name 

163 for d in self.task_node.dimensions 

164 if d.name in self.dataset_type_node.dimensions or not (d.spatial or d.temporal) 

165 } 

166 ) 

167 

168 edge: ReadEdge 

169 """The `~pipeline_graph.PipelineGraph` edge that represents the 

170 prerequisite input connection. 

171 """ 

172 

173 dataset_type_node: DatasetTypeNode 

174 """The `~pipeline_graph.PipelineGraph` node that represents the dataset 

175 type of this connection. 

176 

177 This always uses the registry storage class and is never a component 

178 dataset type. 

179 """ 

180 

181 lookup_function: Callable[ 

182 [DatasetType, Registry, DataCoordinate, Sequence[str]], Iterable[DatasetRef] 

183 ] | None 

184 """A task-provided callback for finding these datasets. 

185 

186 If this is not `None`, it must be used to ensure correct behavior. 

187 """ 

188 

189 dataset_skypix: dict[str, SkyPixDimension] 

190 """Dimensions representing a pixelization of the sky used by the dataset 

191 type for this connection that are also not part of the task's dimensions. 

192 

193 Keys are dimension names. It is at least extremely rare for this 

194 dictionary to have more than one element. 

195 """ 

196 

197 dataset_other_spatial: dict[str, DimensionElement] 

198 """Spatial dimensions other than sky pixelizations used by the dataset type 

199 for this connection that are also not part of the task's dimensions. 

200 """ 

201 

202 dataset_has_timespan: bool 

203 """Whether the dataset has a timespan that should be used in the lookup, 

204 either because it is a calibration dataset or because it has temporal 

205 dimensions that are not part of the tasks's dimensions. 

206 """ 

207 

208 @property 

209 def task_node(self) -> TaskNode: 

210 """The `~pipeline_graph.PipelineGraph` node that represents the task 

211 for this connection. 

212 """ 

213 return self._bounds.task_node 

214 

215 def find( 

216 self, 

217 butler: Butler, 

218 input_collections: Sequence[str], 

219 data_id: DataCoordinate, 

220 skypix_bounds: Mapping[str, RangeSet], 

221 timespan: Timespan | None, 

222 ) -> list[DatasetRef]: 

223 """Find prerequisite input datasets for a single quantum. 

224 

225 Parameters 

226 ---------- 

227 butler : `lsst.daf.butler.Butler` 

228 Butler client to use for queries. 

229 input_collections : `~collections.abc.Sequence` [ `str` ] 

230 Sequence of collections to search, in order. 

231 data_id : `lsst.daf.butler.DataCoordinate` 

232 Data ID for the quantum. 

233 skypix_bounds : `Mapping` [ `str`, `lsst.sphgeom.RangeSet` ] 

234 The spatial bounds of this quantum in various skypix dimensions. 

235 Keys are skypix dimension names (a superset of those in 

236 `dataset_skypix`) and values are sets of integer pixel ID ranges. 

237 timespan : `lsst.daf.butler.Timespan` or `None` 

238 The temporal bounds of this quantum. Guaranteed to not be `None` 

239 if `dataset_has_timespan` is `True`. 

240 

241 Returns 

242 ------- 

243 refs : `list` [ `lsst.daf.butler.DatasetRef` ] 

244 Dataset references. These use 

245 ``self.dataset_type_node.dataset_type``, which may differ from the 

246 connection's dataset type in storage class or [lack of] component. 

247 

248 Raises 

249 ------ 

250 NotImplementedError 

251 Raised for certain relationships between task and dataset type 

252 dimensions that are possible to define but not believed to be 

253 useful in practice. These errors occur late rather than early in 

254 order to allow a `QuantumGraphBuilder` subclass to handle them 

255 first, in case an unusual task's needs must be met by a custom 

256 builder class anyway. 

257 """ 

258 if self.lookup_function: 

259 # If there is a lookup function, just use it; nothing else matters. 

260 return [ 

261 self.dataset_type_node.generalize_ref(ref) 

262 for ref in self.lookup_function( 

263 self.edge.adapt_dataset_type(self.dataset_type_node.dataset_type), 

264 butler.registry, 

265 data_id, 

266 input_collections, 

267 ) 

268 if ref is not None 

269 ] 

270 if self.dataset_type_node.is_calibration: 

271 if self.dataset_type_node.dimensions <= self.constraint_dimensions: 

272 # If this is a calibration dataset and the dataset doesn't have 

273 # any dimensions that aren't constrained by the quantum data 

274 # ID, we know there'll only be one result, and that means we 

275 # can call Registry.findDataset, which takes a timespan. Note 

276 # that the AllDimensionsQuantumGraphBuilder subclass will 

277 # intercept this case in order to optimize it when: 

278 # 

279 # - PipelineTaskConnections.getTemporalBoundsConnections is 

280 # empty; 

281 # 

282 # - the quantum data IDs have temporal dimensions; 

283 # 

284 # and when that happens PrerequisiteFinder.find never gets 

285 # called. 

286 try: 

287 ref = butler.find_dataset( 

288 self.dataset_type_node.dataset_type, 

289 data_id.subset(self.constraint_dimensions), 

290 collections=input_collections, 

291 timespan=timespan, 

292 ) 

293 except MissingDatasetTypeError: 

294 ref = None 

295 return [ref] if ref is not None else [] 

296 else: 

297 extra_dimensions = ( 

298 self.dataset_type_node.dimensions.dimensions - self.constraint_dimensions.dimensions 

299 ) 

300 raise NotImplementedError( 

301 f"No support for calibration lookup {self.task_node.label}.{self.edge.connection_name} " 

302 f"with dimension(s) {extra_dimensions} not fully constrained by the task. " 

303 "Please create a feature-request ticket and use a lookup function in the meantime." 

304 ) 

305 if self.dataset_skypix: 

306 if not self.dataset_has_timespan and not self.dataset_other_spatial: 

307 # If the dataset has skypix dimensions but is not otherwise 

308 # spatial or temporal (this describes reference catalogs and 

309 # things like them), we can stuff the skypix IDs we want into 

310 # the query via bind parameters and call queryDatasets. Once 

311 # again AllDimensionsQuantumGraphBuilder will often intercept 

312 # this case in order to optimize it, when: 

313 # 

314 # - PipelineTaskConnections.getSpatialBoundsConnections is 

315 # empty; 

316 # 

317 # - the quantum data IDs have spatial dimensions; 

318 # 

319 # and when that happens PrerequisiteFinder.find never gets 

320 # called. 

321 where_terms: list[str] = [] 

322 bind: dict[str, list[int]] = {} 

323 for name in self.dataset_skypix: 

324 where_terms.append(f"{name} IN ({name}_pixels)") 

325 pixels: list[int] = [] 

326 for begin, end in skypix_bounds[name]: 

327 pixels.extend(range(begin, end)) 

328 bind[f"{name}_pixels"] = pixels 

329 try: 

330 return list( 

331 butler.registry.queryDatasets( 

332 self.dataset_type_node.dataset_type, 

333 collections=input_collections, 

334 dataId=data_id.subset(self.constraint_dimensions), 

335 where=" AND ".join(where_terms), 

336 bind=bind, 

337 findFirst=True, 

338 ).expanded() 

339 ) 

340 except MissingDatasetTypeError: 

341 return [] 

342 else: 

343 raise NotImplementedError( 

344 f"No support for skypix lookup {self.task_node.label}.{self.edge.connection_name} " 

345 "that requires additional spatial and/or temporal constraints. " 

346 "Please create a feature-request ticket and use a lookup function in the meantime." 

347 ) 

348 if self._bounds.spatial_connections or self._bounds.temporal_connections: 

349 raise NotImplementedError( 

350 f"No support for prerequisite lookup {self.task_node.label}.{self.edge.connection_name} " 

351 "that requires other connections to determine spatial or temporal bounds but does not " 

352 "fit into one of our standard cases. " 

353 "Please create a feature-request ticket and use a lookup function in the meantime." 

354 ) 

355 # If the spatial/temporal bounds are not customized, and the dataset 

356 # doesn't have any skypix dimensions, a vanilla queryDatasets call 

357 # should work. This case should always be optimized by 

358 # AllDimensionsQuantumGraphBuilder as well. Note that we use the 

359 # original quantum data ID here, not those with constraint_dimensions 

360 # that strips out the spatial/temporal stuff, because here we want the 

361 # butler query system to handle the spatial/temporal stuff like it 

362 # normally would. 

363 try: 

364 return list( 

365 butler.registry.queryDatasets( 

366 self.dataset_type_node.dataset_type, 

367 collections=input_collections, 

368 dataId=data_id, 

369 findFirst=True, 

370 ).expanded() 

371 ) 

372 except MissingDatasetTypeError: 

373 return [] 

374 

375 

376@dataclasses.dataclass 

377class PrerequisiteBounds: 

378 """A QuantumGraph-generation helper class that manages the spatial and 

379 temporal bounds of a tasks' quanta, for the purpose of finding 

380 prerequisite inputs. 

381 """ 

382 

383 task_node: TaskNode 

384 """The `~pipeline_graph.PipelineGraph` node that represents the task.""" 

385 

386 spatial_connections: frozenset[str] = dataclasses.field(init=False) 

387 """Regular input or output connections whose (assumed spatial) data IDs 

388 should be used to define the spatial bounds of this task's quanta. 

389 

390 See Also 

391 -------- 

392 PipelineTaskConnections.getSpatialBoundsConnections 

393 """ 

394 

395 temporal_connections: frozenset[str] = dataclasses.field(init=False) 

396 """Regular input or output connections whose (assumed temporal) data IDs 

397 should be used to define the temporal bounds of this task's quanta. 

398 

399 See Also 

400 -------- 

401 PipelineTaskConnections.getTemporalBoundsConnections 

402 """ 

403 

404 all_dataset_skypix: dict[str, SkyPixDimension] = dataclasses.field(default_factory=dict) 

405 """The union of all `PrerequisiteFinder.dataset_skypix` attributes for all 

406 (remaining) prerequisite finders for this task. 

407 """ 

408 

409 any_dataset_has_timespan: bool = dataclasses.field(default=False) 

410 """Whether any `PrerequisiteFinder.dataset_has_timespan` attribute is true 

411 for any (remaining) prerequisite finder for this task. 

412 """ 

413 

414 def __post_init__(self) -> None: 

415 self.spatial_connections = frozenset(self.task_node.get_spatial_bounds_connections()) 

416 self.temporal_connections = frozenset(self.task_node.get_temporal_bounds_connections()) 

417 

418 def make_skypix_bounds_builder(self, quantum_data_id: DataCoordinate) -> SkyPixBoundsBuilder: 

419 """Return an object that accumulates the appropriate spatial bounds for 

420 a quantum. 

421 

422 Parameters 

423 ---------- 

424 quantum_data_id : `lsst.daf.butler.DataCoordinate` 

425 Data ID for this quantum. 

426 

427 Returns 

428 ------- 

429 builder : `SkyPixBoundsBuilder` 

430 Object that accumulates the appropriate spatial bounds for a 

431 quantum. If the spatial bounds are not needed, this object will do 

432 nothing. 

433 """ 

434 if not self.all_dataset_skypix: 

435 return _TrivialSkyPixBoundsBuilder() 

436 if self.spatial_connections: 

437 return _ConnectionSkyPixBoundsBuilder( 

438 self.task_node, self.spatial_connections, self.all_dataset_skypix.values(), quantum_data_id 

439 ) 

440 if self.task_node.dimensions.spatial: 

441 return _QuantumOnlySkyPixBoundsBuilder(self.all_dataset_skypix.values(), quantum_data_id) 

442 else: 

443 return _UnboundedSkyPixBoundsBuilder(self.all_dataset_skypix.values()) 

444 

445 def make_timespan_builder(self, quantum_data_id: DataCoordinate) -> TimespanBuilder: 

446 """Return an object that accumulates the appropriate timespan for 

447 a quantum. 

448 

449 Parameters 

450 ---------- 

451 quantum_data_id : `lsst.daf.butler.DataCoordinate` 

452 Data ID for this quantum. 

453 

454 Returns 

455 ------- 

456 builder : `TimespanBuilder` 

457 Object that accumulates the appropriate timespan bounds for a 

458 quantum. If a timespan is not needed, this object will do nothing. 

459 """ 

460 if not self.any_dataset_has_timespan: 

461 return _TrivialTimespanBuilder() 

462 if self.temporal_connections: 

463 return _ConnectionTimespanBuilder(self.task_node, self.temporal_connections, quantum_data_id) 

464 if self.task_node.dimensions.temporal: 

465 return _QuantumOnlyTimespanBuilder(quantum_data_id) 

466 else: 

467 return _UnboundedTimespanBuilder() 

468 

469 

470class SkyPixBoundsBuilder(ABC): 

471 """A base class for objects that accumulate the appropriate spatial bounds 

472 for a quantum. 

473 """ 

474 

475 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

476 """Handle the skeleton graph node for a regular input/output connection 

477 for this quantum, including its data ID in the bounds if appropriate. 

478 

479 Parameters 

480 ---------- 

481 parent_dataset_type_name : `str` 

482 Name of the dataset type. Never a component dataset type name. 

483 data_id : `lsst.daf.butler.DataCoordinate` 

484 Data ID for the dataset. 

485 """ 

486 pass 

487 

488 @abstractmethod 

489 def finish(self) -> dict[str, RangeSet]: 

490 """Finish building the spatial bounds and return them. 

491 

492 Returns 

493 ------- 

494 bounds : `dict` [ `str`, `lsst.sphgeom.RangeSet` ] 

495 The spatial bounds of this quantum in various skypix dimensions. 

496 Keys are skypix dimension names and values are sets of integer 

497 pixel ID ranges. 

498 """ 

499 raise NotImplementedError() 

500 

501 

502class TimespanBuilder(ABC): 

503 """A base class for objects that accumulate the appropriate timespan 

504 for a quantum. 

505 """ 

506 

507 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

508 """Handle the skeleton graph node for a regular input/output connection 

509 for this quantum, including its data ID in the bounds if appropriate. 

510 

511 Parameters 

512 ---------- 

513 parent_dataset_type_name : `str` 

514 Name of the dataset type. Never a component dataset type name. 

515 data_id : `lsst.daf.butler.DataCoordinate` 

516 Data ID for the dataset. 

517 """ 

518 pass 

519 

520 @abstractmethod 

521 def finish(self) -> Timespan | None: 

522 """Finish building the timespan and return it. 

523 

524 Returns 

525 ------- 

526 timespan : `lsst.daf.butler.Timespan` or `None` 

527 The timespan of this quantum, or `None` if it is known to not be 

528 needed. 

529 """ 

530 raise NotImplementedError() 

531 

532 

533class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

534 """Implementation of `SkyPixBoundsBuilder` for when no skypix bounds are 

535 needed. 

536 """ 

537 

538 def finish(self) -> dict[str, RangeSet]: 

539 return {} 

540 

541 

542class _TrivialTimespanBuilder(TimespanBuilder): 

543 """Implementation of `TimespanBuilder` for when no timespan is needed.""" 

544 

545 def finish(self) -> None: 

546 return None 

547 

548 

549class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder): 

550 """Implementation of `SkyPixBoundsBuilder` for when the quantum data IDs 

551 provide the only relevant spatial regions. 

552 """ 

553 

554 def __init__(self, dimensions: Iterable[SkyPixDimension], quantum_data_id: DataCoordinate) -> None: 

555 self._region = quantum_data_id.region 

556 self._dimensions = dimensions 

557 

558 def finish(self) -> dict[str, RangeSet]: 

559 return { 

560 dimension.name: dimension.pixelization.envelope(self._region) for dimension in self._dimensions 

561 } 

562 

563 

564class _QuantumOnlyTimespanBuilder(TimespanBuilder): 

565 """Implementation of `TimespanBuilder` for when the quantum data IDs 

566 provide the only relevant timespans. 

567 """ 

568 

569 def __init__(self, quantum_data_id: DataCoordinate) -> None: 

570 self._timespan = cast(Timespan, quantum_data_id.timespan) 

571 

572 def finish(self) -> Timespan: 

573 return self._timespan 

574 

575 

576class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

577 """Implementation of `SkyPixBoundsBuilder` for when the bounds cover the 

578 full sky. 

579 """ 

580 

581 def __init__(self, dimensions: Iterable[SkyPixDimension]): 

582 self._dimensions = dimensions 

583 

584 def finish(self) -> dict[str, RangeSet]: 

585 return {dimension.name: dimension.pixelization.universe() for dimension in self._dimensions} 

586 

587 

588class _UnboundedTimespanBuilder(TimespanBuilder): 

589 """Implementation of `TimespanBuilder` for when the timespan covers all 

590 time. 

591 """ 

592 

593 def finish(self) -> Timespan: 

594 return Timespan(None, None) 

595 

596 

597class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

598 """Implementation of `SkyPixBoundsBuilder` for when other input or output 

599 connections contribute to the spatial bounds. 

600 """ 

601 

602 def __init__( 

603 self, 

604 task_node: TaskNode, 

605 bounds_connections: frozenset[str], 

606 dimensions: Iterable[SkyPixDimension], 

607 quantum_data_id: DataCoordinate, 

608 ) -> None: 

609 self._dimensions = dimensions 

610 self._regions: list[Region] = [] 

611 if task_node.dimensions.spatial: 

612 self._regions.append(quantum_data_id.region) 

613 self._dataset_type_names: set[str] = set() 

614 for connection_name in bounds_connections: 

615 if edge := task_node.inputs.get(connection_name): 

616 self._dataset_type_names.add(edge.parent_dataset_type_name) 

617 else: 

618 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name) 

619 # Note that we end up raising if the input is a prerequisite (and 

620 # hence not in task_node.inputs or task_node.outputs); this 

621 # justifies the cast in `handle_dataset`. 

622 

623 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

624 if parent_dataset_type_name in self._dataset_type_names: 

625 self._regions.append(data_id.region) 

626 

627 def finish(self) -> dict[str, RangeSet]: 

628 result = {} 

629 for dimension in self._dimensions: 

630 bounds = RangeSet() 

631 for region in self._regions: 

632 bounds |= dimension.pixelization.envelope(region) 

633 result[dimension.name] = bounds 

634 return result 

635 

636 

637class _ConnectionTimespanBuilder(TimespanBuilder): 

638 """Implementation of `TimespanBuilder` for when other input or output 

639 connections contribute to the timespan. 

640 """ 

641 

642 def __init__( 

643 self, 

644 task_node: TaskNode, 

645 bounds_connections: frozenset[str], 

646 quantum_data_id: DataCoordinate, 

647 ) -> None: 

648 timespan = ( 

649 cast(Timespan, quantum_data_id.timespan) 

650 if task_node.dimensions.temporal 

651 else Timespan.makeEmpty() 

652 ) 

653 self._begin_nsec = timespan._nsec[0] 

654 self._end_nsec = timespan._nsec[1] 

655 self._dataset_type_names = set() 

656 for connection_name in bounds_connections: 

657 if edge := task_node.inputs.get(connection_name): 

658 self._dataset_type_names.add(edge.parent_dataset_type_name) 

659 else: 

660 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name) 

661 # Note that we end up raising if the input is a prerequisite (and 

662 # hence not in task_node.inputs or task_node.outputs); this 

663 # justifies the cast in `handle_dataset`. 

664 

665 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

666 if parent_dataset_type_name in self._dataset_type_names: 

667 nsec = cast(Timespan, data_id.timespan)._nsec 

668 self._begin_nsec = min(self._begin_nsec, nsec[0]) 

669 self._end_nsec = max(self._end_nsec, nsec[1]) 

670 

671 def finish(self) -> Timespan: 

672 return Timespan(None, None, _nsec=(self._begin_nsec, self._end_nsec))