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

202 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-11-30 12:09 +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( 

151 self.dataset_type_node.dimensions.elements.names, 

152 self.dataset_type_node.dimensions.universe, 

153 ) 

154 if isinstance(best_spatial_element, SkyPixDimension): 

155 self.dataset_skypix[best_spatial_element.name] = best_spatial_element 

156 else: 

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

158 DimensionElement, best_spatial_element 

159 ) 

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

161 self.dataset_type_node.dimensions.temporal - self.task_node.dimensions.temporal 

162 ) 

163 new_constraint_dimensions = set() 

164 universe = self.task_node.dimensions.universe 

165 for dimension_name in self.task_node.dimensions.names: 

166 if dimension_name in self.dataset_type_node.dimensions.names: 

167 new_constraint_dimensions.add(dimension_name) 

168 else: 

169 dimension = universe[dimension_name] 

170 if not (dimension.spatial or dimension.temporal): 

171 new_constraint_dimensions.add(dimension_name) 

172 self.constraint_dimensions = universe.conform(new_constraint_dimensions) 

173 

174 edge: ReadEdge 

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

176 prerequisite input connection. 

177 """ 

178 

179 dataset_type_node: DatasetTypeNode 

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

181 type of this connection. 

182 

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

184 dataset type. 

185 """ 

186 

187 lookup_function: Callable[ 

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

189 ] | None 

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

191 

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

193 """ 

194 

195 dataset_skypix: dict[str, SkyPixDimension] 

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

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

198 

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

200 dictionary to have more than one element. 

201 """ 

202 

203 dataset_other_spatial: dict[str, DimensionElement] 

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

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

206 """ 

207 

208 dataset_has_timespan: bool 

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

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

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

212 """ 

213 

214 @property 

215 def task_node(self) -> TaskNode: 

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

217 for this connection. 

218 """ 

219 return self._bounds.task_node 

220 

221 def find( 

222 self, 

223 butler: Butler, 

224 input_collections: Sequence[str], 

225 data_id: DataCoordinate, 

226 skypix_bounds: Mapping[str, RangeSet], 

227 timespan: Timespan | None, 

228 ) -> list[DatasetRef]: 

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

230 

231 Parameters 

232 ---------- 

233 butler : `lsst.daf.butler.Butler` 

234 Butler client to use for queries. 

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

236 Sequence of collections to search, in order. 

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

238 Data ID for the quantum. 

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

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

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

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

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

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

245 if `dataset_has_timespan` is `True`. 

246 

247 Returns 

248 ------- 

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

250 Dataset references. These use 

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

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

253 

254 Raises 

255 ------ 

256 NotImplementedError 

257 Raised for certain relationships between task and dataset type 

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

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

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

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

262 builder class anyway. 

263 """ 

264 if self.lookup_function: 

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

266 return [ 

267 self.dataset_type_node.generalize_ref(ref) 

268 for ref in self.lookup_function( 

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

270 butler.registry, 

271 data_id, 

272 input_collections, 

273 ) 

274 if ref is not None 

275 ] 

276 if self.dataset_type_node.is_calibration: 

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

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

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

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

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

282 # that the AllDimensionsQuantumGraphBuilder subclass will 

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

284 # 

285 # - PipelineTaskConnections.getTemporalBoundsConnections is 

286 # empty; 

287 # 

288 # - the quantum data IDs have temporal dimensions; 

289 # 

290 # and when that happens PrerequisiteFinder.find never gets 

291 # called. 

292 try: 

293 ref = butler.find_dataset( 

294 self.dataset_type_node.dataset_type, 

295 data_id.subset(self.constraint_dimensions), 

296 collections=input_collections, 

297 timespan=timespan, 

298 ) 

299 except MissingDatasetTypeError: 

300 ref = None 

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

302 else: 

303 extra_dimensions = self.dataset_type_node.dimensions.names - self.constraint_dimensions.names 

304 raise NotImplementedError( 

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

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

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

308 ) 

309 if self.dataset_skypix: 

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

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

312 # spatial or temporal (this describes reference catalogs and 

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

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

315 # again AllDimensionsQuantumGraphBuilder will often intercept 

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

317 # 

318 # - PipelineTaskConnections.getSpatialBoundsConnections is 

319 # empty; 

320 # 

321 # - the quantum data IDs have spatial dimensions; 

322 # 

323 # and when that happens PrerequisiteFinder.find never gets 

324 # called. 

325 where_terms: list[str] = [] 

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

327 for name in self.dataset_skypix: 

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

329 pixels: list[int] = [] 

330 for begin, end in skypix_bounds[name]: 

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

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

333 try: 

334 return list( 

335 butler.registry.queryDatasets( 

336 self.dataset_type_node.dataset_type, 

337 collections=input_collections, 

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

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

340 bind=bind, 

341 findFirst=True, 

342 ).expanded() 

343 ) 

344 except MissingDatasetTypeError: 

345 return [] 

346 else: 

347 raise NotImplementedError( 

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

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

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

351 ) 

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

353 raise NotImplementedError( 

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

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

356 "fit into one of our standard cases. " 

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

358 ) 

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

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

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

362 # AllDimensionsQuantumGraphBuilder as well. Note that we use the 

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

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

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

366 # normally would. 

367 try: 

368 return list( 

369 butler.registry.queryDatasets( 

370 self.dataset_type_node.dataset_type, 

371 collections=input_collections, 

372 dataId=data_id, 

373 findFirst=True, 

374 ).expanded() 

375 ) 

376 except MissingDatasetTypeError: 

377 return [] 

378 

379 

380@dataclasses.dataclass 

381class PrerequisiteBounds: 

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

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

384 prerequisite inputs. 

385 """ 

386 

387 task_node: TaskNode 

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

389 

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

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

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

393 

394 See Also 

395 -------- 

396 PipelineTaskConnections.getSpatialBoundsConnections 

397 """ 

398 

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

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

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

402 

403 See Also 

404 -------- 

405 PipelineTaskConnections.getTemporalBoundsConnections 

406 """ 

407 

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

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

410 (remaining) prerequisite finders for this task. 

411 """ 

412 

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

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

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

416 """ 

417 

418 def __post_init__(self) -> None: 

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

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

421 

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

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

424 a quantum. 

425 

426 Parameters 

427 ---------- 

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

429 Data ID for this quantum. 

430 

431 Returns 

432 ------- 

433 builder : `SkyPixBoundsBuilder` 

434 Object that accumulates the appropriate spatial bounds for a 

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

436 nothing. 

437 """ 

438 if not self.all_dataset_skypix: 

439 return _TrivialSkyPixBoundsBuilder() 

440 if self.spatial_connections: 

441 return _ConnectionSkyPixBoundsBuilder( 

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

443 ) 

444 if self.task_node.dimensions.spatial: 

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

446 else: 

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

448 

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

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

451 a quantum. 

452 

453 Parameters 

454 ---------- 

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

456 Data ID for this quantum. 

457 

458 Returns 

459 ------- 

460 builder : `TimespanBuilder` 

461 Object that accumulates the appropriate timespan bounds for a 

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

463 """ 

464 if not self.any_dataset_has_timespan: 

465 return _TrivialTimespanBuilder() 

466 if self.temporal_connections: 

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

468 if self.task_node.dimensions.temporal: 

469 return _QuantumOnlyTimespanBuilder(quantum_data_id) 

470 else: 

471 return _UnboundedTimespanBuilder() 

472 

473 

474class SkyPixBoundsBuilder(ABC): 

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

476 for a quantum. 

477 """ 

478 

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

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

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

482 

483 Parameters 

484 ---------- 

485 parent_dataset_type_name : `str` 

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

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

488 Data ID for the dataset. 

489 """ 

490 pass 

491 

492 @abstractmethod 

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

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

495 

496 Returns 

497 ------- 

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

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

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

501 pixel ID ranges. 

502 """ 

503 raise NotImplementedError() 

504 

505 

506class TimespanBuilder(ABC): 

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

508 for a quantum. 

509 """ 

510 

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

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

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

514 

515 Parameters 

516 ---------- 

517 parent_dataset_type_name : `str` 

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

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

520 Data ID for the dataset. 

521 """ 

522 pass 

523 

524 @abstractmethod 

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

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

527 

528 Returns 

529 ------- 

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

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

532 needed. 

533 """ 

534 raise NotImplementedError() 

535 

536 

537class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

539 needed. 

540 """ 

541 

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

543 return {} 

544 

545 

546class _TrivialTimespanBuilder(TimespanBuilder): 

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

548 

549 def finish(self) -> None: 

550 return None 

551 

552 

553class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

555 provide the only relevant spatial regions. 

556 """ 

557 

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

559 self._region = quantum_data_id.region 

560 self._dimensions = dimensions 

561 

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

563 return { 

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

565 } 

566 

567 

568class _QuantumOnlyTimespanBuilder(TimespanBuilder): 

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

570 provide the only relevant timespans. 

571 """ 

572 

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

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

575 

576 def finish(self) -> Timespan: 

577 return self._timespan 

578 

579 

580class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

582 full sky. 

583 """ 

584 

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

586 self._dimensions = dimensions 

587 

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

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

590 

591 

592class _UnboundedTimespanBuilder(TimespanBuilder): 

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

594 time. 

595 """ 

596 

597 def finish(self) -> Timespan: 

598 return Timespan(None, None) 

599 

600 

601class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

603 connections contribute to the spatial bounds. 

604 """ 

605 

606 def __init__( 

607 self, 

608 task_node: TaskNode, 

609 bounds_connections: frozenset[str], 

610 dimensions: Iterable[SkyPixDimension], 

611 quantum_data_id: DataCoordinate, 

612 ) -> None: 

613 self._dimensions = dimensions 

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

615 if task_node.dimensions.spatial: 

616 self._regions.append(quantum_data_id.region) 

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

618 for connection_name in bounds_connections: 

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

620 self._dataset_type_names.add(edge.parent_dataset_type_name) 

621 else: 

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

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

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

625 # justifies the cast in `handle_dataset`. 

626 

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

628 if parent_dataset_type_name in self._dataset_type_names: 

629 self._regions.append(data_id.region) 

630 

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

632 result = {} 

633 for dimension in self._dimensions: 

634 bounds = RangeSet() 

635 for region in self._regions: 

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

637 result[dimension.name] = bounds 

638 return result 

639 

640 

641class _ConnectionTimespanBuilder(TimespanBuilder): 

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

643 connections contribute to the timespan. 

644 """ 

645 

646 def __init__( 

647 self, 

648 task_node: TaskNode, 

649 bounds_connections: frozenset[str], 

650 quantum_data_id: DataCoordinate, 

651 ) -> None: 

652 timespan = ( 

653 cast(Timespan, quantum_data_id.timespan) 

654 if task_node.dimensions.temporal 

655 else Timespan.makeEmpty() 

656 ) 

657 self._begin_nsec = timespan._nsec[0] 

658 self._end_nsec = timespan._nsec[1] 

659 self._dataset_type_names = set() 

660 for connection_name in bounds_connections: 

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

662 self._dataset_type_names.add(edge.parent_dataset_type_name) 

663 else: 

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

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

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

667 # justifies the cast in `handle_dataset`. 

668 

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

670 if parent_dataset_type_name in self._dataset_type_names: 

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

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

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

674 

675 def finish(self) -> Timespan: 

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