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

194 statements  

« prev     ^ index     » next       coverage.py v7.3.0, created at 2023-08-31 09:39 +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 program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

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

23QuantumGraph generation. 

24""" 

25 

26from __future__ import annotations 

27 

28__all__ = ( 

29 "SkyPixBoundsBuilder", 

30 "PrerequisiteFinder", 

31 "PrerequisiteBounds", 

32 "TimespanBuilder", 

33 "PrerequisiteInfo", 

34) 

35 

36import dataclasses 

37from abc import ABC, abstractmethod 

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

39from typing import cast 

40 

41from lsst.daf.butler import ( 

42 Butler, 

43 DataCoordinate, 

44 DatasetRef, 

45 DatasetType, 

46 DimensionElement, 

47 Registry, 

48 SkyPixDimension, 

49 Timespan, 

50) 

51from lsst.daf.butler.registry import MissingDatasetTypeError 

52from lsst.sphgeom import RangeSet, Region 

53 

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

55 

56 

57@dataclasses.dataclass 

58class PrerequisiteInfo: 

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

60 prerequisite input connections for a task. 

61 """ 

62 

63 bounds: PrerequisiteBounds 

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

65 task's quanta. 

66 """ 

67 

68 finders: dict[str, PrerequisiteFinder] 

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

70 connection. 

71 

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

73 by implementations of `QuantumGraphBuilder.process_subgraph` to take 

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

75 base class. 

76 """ 

77 

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

79 self.bounds = PrerequisiteBounds(task_node) 

80 self.finders = { 

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

82 for edge in task_node.prerequisite_inputs.values() 

83 } 

84 

85 def update_bounds(self) -> None: 

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

87 reflect the needs of only the finders that remain. 

88 """ 

89 self.bounds.all_dataset_skypix.clear() 

90 self.bounds.any_dataset_has_timespan = False 

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

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

93 self.bounds.any_dataset_has_timespan = ( 

94 self.bounds.any_dataset_has_timespan or finder.dataset_has_timespan 

95 ) 

96 

97 

98class PrerequisiteFinder: 

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

100 prerequisite input connection. 

101 

102 Parameters 

103 ---------- 

104 edge : `pipeline_graph.ReadEdge` 

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

106 prerequisite input connection. 

107 bounds : `PrerequisiteBounds` 

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

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

110 pipeline_graph `pipeline_graph.PipelineGraph` 

111 Graph representation of the pipeline. 

112 

113 Notes 

114 ----- 

115 `PrerequisiteFinder` instances are usually constructed by a 

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

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

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

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

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

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

122 removing those finder instances after those prerequisites have been found 

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

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

125 later in `QuantumGraphBuilder.build`. 

126 """ 

127 

128 def __init__( 

129 self, 

130 edge: ReadEdge, 

131 bounds: PrerequisiteBounds, 

132 pipeline_graph: PipelineGraph, 

133 ): 

134 self.edge = edge 

135 self._bounds = bounds 

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

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

138 self.dataset_skypix = {} 

139 self.dataset_other_spatial = {} 

140 self.dataset_has_timespan = False 

141 self.constraint_dimensions = self.task_node.dimensions 

142 if self.lookup_function is None: 

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

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

145 if isinstance(best_spatial_element, SkyPixDimension): 

146 self.dataset_skypix[best_spatial_element.name] = best_spatial_element 

147 else: 

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

149 DimensionElement, best_spatial_element 

150 ) 

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

152 self.dataset_type_node.dimensions.temporal - self.task_node.dimensions.temporal 

153 ) 

154 self.constraint_dimensions = self.constraint_dimensions.universe.extract( 

155 { 

156 d.name 

157 for d in self.task_node.dimensions 

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

159 } 

160 ) 

161 

162 edge: ReadEdge 

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

164 prerequisite input connection. 

165 """ 

166 

167 dataset_type_node: DatasetTypeNode 

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

169 type of this connection. 

170 

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

172 dataset type. 

173 """ 

174 

175 lookup_function: Callable[ 

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

177 ] | None 

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

179 

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

181 """ 

182 

183 dataset_skypix: dict[str, SkyPixDimension] 

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

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

186 

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

188 dictionary to have more than one element. 

189 """ 

190 

191 dataset_other_spatial: dict[str, DimensionElement] 

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

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

194 """ 

195 

196 dataset_has_timespan: bool 

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

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

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

200 """ 

201 

202 @property 

203 def task_node(self) -> TaskNode: 

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

205 for this connection. 

206 """ 

207 return self._bounds.task_node 

208 

209 def find( 

210 self, 

211 butler: Butler, 

212 input_collections: Sequence[str], 

213 data_id: DataCoordinate, 

214 skypix_bounds: Mapping[str, RangeSet], 

215 timespan: Timespan | None, 

216 ) -> list[DatasetRef]: 

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

218 

219 Parameters 

220 ---------- 

221 butler : `lsst.daf.butler.Butler` 

222 Butler client to use for queries. 

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

224 Sequence of collections to search, in order. 

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

226 Data ID for the quantum. 

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

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

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

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

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

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

233 if `dataset_has_timespan` is `True`. 

234 

235 Returns 

236 ------- 

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

238 Dataset references. These use 

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

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

241 

242 Raises 

243 ------ 

244 NotImplementedError 

245 Raised for certain relationships between task and dataset type 

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

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

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

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

250 builder class anyway. 

251 """ 

252 if self.lookup_function: 

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

254 return [ 

255 self.dataset_type_node.generalize_ref(ref) 

256 for ref in self.lookup_function( 

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

258 butler.registry, 

259 data_id, 

260 input_collections, 

261 ) 

262 if ref is not None 

263 ] 

264 if self.dataset_type_node.is_calibration: 

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

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

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

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

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

270 # that the AllDimensionsQuantumGraphBuilder subclass will 

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

272 # 

273 # - PipelineTaskConnections.getTemporalBoundsConnections is 

274 # empty; 

275 # 

276 # - the quantum data IDs have temporal dimensions; 

277 # 

278 # and when that happens PrerequisiteFinder.find never gets 

279 # called. 

280 try: 

281 ref = butler.registry.findDataset( 

282 self.dataset_type_node.dataset_type, 

283 data_id.subset(self.constraint_dimensions), 

284 collections=input_collections, 

285 timespan=timespan, 

286 ) 

287 except MissingDatasetTypeError: 

288 ref = None 

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

290 else: 

291 extra_dimensions = ( 

292 self.dataset_type_node.dimensions.dimensions - self.constraint_dimensions.dimensions 

293 ) 

294 raise NotImplementedError( 

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

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

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

298 ) 

299 if self.dataset_skypix: 

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

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

302 # spatial or temporal (this describes reference catalogs and 

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

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

305 # again AllDimensionsQuantumGraphBuilder will often intercept 

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

307 # 

308 # - PipelineTaskConnections.getSpatialBoundsConnections is 

309 # empty; 

310 # 

311 # - the quantum data IDs have spatial dimensions; 

312 # 

313 # and when that happens PrerequisiteFinder.find never gets 

314 # called. 

315 where_terms: list[str] = [] 

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

317 for name in self.dataset_skypix: 

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

319 pixels: list[int] = [] 

320 for begin, end in skypix_bounds[name]: 

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

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

323 try: 

324 return list( 

325 butler.registry.queryDatasets( 

326 self.dataset_type_node.dataset_type, 

327 collections=input_collections, 

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

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

330 bind=bind, 

331 findFirst=True, 

332 ).expanded() 

333 ) 

334 except MissingDatasetTypeError: 

335 return [] 

336 else: 

337 raise NotImplementedError( 

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

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

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

341 ) 

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

343 raise NotImplementedError( 

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

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

346 "fit into one of our standard cases. " 

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

348 ) 

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

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

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

352 # AllDimensionsQuantumGraphBuilder as well. Note that we use the 

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

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

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

356 # normally would. 

357 try: 

358 return list( 

359 butler.registry.queryDatasets( 

360 self.dataset_type_node.dataset_type, 

361 collections=input_collections, 

362 dataId=data_id, 

363 findFirst=True, 

364 ).expanded() 

365 ) 

366 except MissingDatasetTypeError: 

367 return [] 

368 

369 

370@dataclasses.dataclass 

371class PrerequisiteBounds: 

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

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

374 prerequisite inputs. 

375 """ 

376 

377 task_node: TaskNode 

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

379 

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

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

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

383 

384 See Also 

385 -------- 

386 PipelineTaskConnections.getSpatialBoundsConnections 

387 """ 

388 

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

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

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

392 

393 See Also 

394 -------- 

395 PipelineTaskConnections.getTemporalBoundsConnections 

396 """ 

397 

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

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

400 (remaining) prerequisite finders for this task. 

401 """ 

402 

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

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

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

406 """ 

407 

408 def __post_init__(self) -> None: 

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

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

411 

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

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

414 a quantum. 

415 

416 Parameters 

417 ---------- 

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

419 Data ID for this quantum. 

420 

421 Returns 

422 ------- 

423 builder : `SkyPixBoundsBuilder` 

424 Object that accumulates the appropriate spatial bounds for a 

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

426 nothing. 

427 """ 

428 if not self.all_dataset_skypix: 

429 return _TrivialSkyPixBoundsBuilder() 

430 if self.spatial_connections: 

431 return _ConnectionSkyPixBoundsBuilder( 

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

433 ) 

434 if self.task_node.dimensions.spatial: 

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

436 else: 

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

438 

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

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

441 a quantum. 

442 

443 Parameters 

444 ---------- 

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

446 Data ID for this quantum. 

447 

448 Returns 

449 ------- 

450 builder : `TimespanBuilder` 

451 Object that accumulates the appropriate timespan bounds for a 

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

453 """ 

454 if not self.any_dataset_has_timespan: 

455 return _TrivialTimespanBuilder() 

456 if self.temporal_connections: 

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

458 if self.task_node.dimensions.temporal: 

459 return _QuantumOnlyTimespanBuilder(quantum_data_id) 

460 else: 

461 return _UnboundedTimespanBuilder() 

462 

463 

464class SkyPixBoundsBuilder(ABC): 

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

466 for a quantum. 

467 """ 

468 

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

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

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

472 

473 Parameters 

474 ---------- 

475 parent_dataset_type_name : `str` 

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

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

478 Data ID for the dataset. 

479 """ 

480 pass 

481 

482 @abstractmethod 

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

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

485 

486 Returns 

487 ------- 

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

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

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

491 pixel ID ranges. 

492 """ 

493 raise NotImplementedError() 

494 

495 

496class TimespanBuilder(ABC): 

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

498 for a quantum. 

499 """ 

500 

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

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

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

504 

505 Parameters 

506 ---------- 

507 parent_dataset_type_name : `str` 

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

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

510 Data ID for the dataset. 

511 """ 

512 pass 

513 

514 @abstractmethod 

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

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

517 

518 Returns 

519 ------- 

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

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

522 needed. 

523 """ 

524 raise NotImplementedError() 

525 

526 

527class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

529 needed. 

530 """ 

531 

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

533 return {} 

534 

535 

536class _TrivialTimespanBuilder(TimespanBuilder): 

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

538 

539 def finish(self) -> None: 

540 return None 

541 

542 

543class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

545 provide the only relevant spatial regions. 

546 """ 

547 

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

549 self._region = quantum_data_id.region 

550 self._dimensions = dimensions 

551 

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

553 return { 

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

555 } 

556 

557 

558class _QuantumOnlyTimespanBuilder(TimespanBuilder): 

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

560 provide the only relevant timespans. 

561 """ 

562 

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

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

565 

566 def finish(self) -> Timespan: 

567 return self._timespan 

568 

569 

570class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

572 full sky. 

573 """ 

574 

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

576 self._dimensions = dimensions 

577 

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

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

580 

581 

582class _UnboundedTimespanBuilder(TimespanBuilder): 

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

584 time. 

585 """ 

586 

587 def finish(self) -> Timespan: 

588 return Timespan(None, None) 

589 

590 

591class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

593 connections contribute to the spatial bounds. 

594 """ 

595 

596 def __init__( 

597 self, 

598 task_node: TaskNode, 

599 bounds_connections: frozenset[str], 

600 dimensions: Iterable[SkyPixDimension], 

601 quantum_data_id: DataCoordinate, 

602 ) -> None: 

603 self._dimensions = dimensions 

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

605 if task_node.dimensions.spatial: 

606 self._regions.append(quantum_data_id.region) 

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

608 for connection_name in bounds_connections: 

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

610 self._dataset_type_names.add(edge.parent_dataset_type_name) 

611 else: 

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

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

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

615 # justifies the cast in `handle_dataset`. 

616 

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

618 if parent_dataset_type_name in self._dataset_type_names: 

619 self._regions.append(data_id.region) 

620 

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

622 result = {} 

623 for dimension in self._dimensions: 

624 bounds = RangeSet() 

625 for region in self._regions: 

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

627 result[dimension.name] = bounds 

628 return result 

629 

630 

631class _ConnectionTimespanBuilder(TimespanBuilder): 

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

633 connections contribute to the timespan. 

634 """ 

635 

636 def __init__( 

637 self, 

638 task_node: TaskNode, 

639 bounds_connections: frozenset[str], 

640 quantum_data_id: DataCoordinate, 

641 ) -> None: 

642 timespan = ( 

643 cast(Timespan, quantum_data_id.timespan) 

644 if task_node.dimensions.temporal 

645 else Timespan.makeEmpty() 

646 ) 

647 self._begin_nsec = timespan._nsec[0] 

648 self._end_nsec = timespan._nsec[1] 

649 self._dataset_type_names = set() 

650 for connection_name in bounds_connections: 

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

652 self._dataset_type_names.add(edge.parent_dataset_type_name) 

653 else: 

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

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

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

657 # justifies the cast in `handle_dataset`. 

658 

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

660 if parent_dataset_type_name in self._dataset_type_names: 

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

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

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

664 

665 def finish(self) -> Timespan: 

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