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

201 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-09 01:22 -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/>. 

27 

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

29QuantumGraph generation. 

30""" 

31 

32from __future__ import annotations 

33 

34__all__ = ( 

35 "PrerequisiteBounds", 

36 "PrerequisiteFinder", 

37 "PrerequisiteInfo", 

38 "SkyPixBoundsBuilder", 

39 "TimespanBuilder", 

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 Parameters 

69 ---------- 

70 task_node : `TaskNode` 

71 The relevant node. 

72 pipeline_graph : `PipelineGraph` 

73 The pipeline graph. 

74 """ 

75 

76 bounds: PrerequisiteBounds 

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

78 task's quanta. 

79 """ 

80 

81 finders: dict[str, PrerequisiteFinder] 

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

83 connection. 

84 

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 """ 

90 

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 } 

97 

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 ) 

109 

110 

111class PrerequisiteFinder: 

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

113 prerequisite input connection. 

114 

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. 

125 

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 """ 

140 

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) 

189 

190 edge: ReadEdge 

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

192 prerequisite input connection. 

193 """ 

194 

195 dataset_type_node: DatasetTypeNode 

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

197 type of this connection. 

198 

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

200 dataset type. 

201 """ 

202 

203 lookup_function: ( 

204 Callable[[DatasetType, Registry, DataCoordinate, Sequence[str]], Iterable[DatasetRef]] | None 

205 ) 

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

207 

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

209 """ 

210 

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. 

214 

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

216 dictionary to have more than one element. 

217 """ 

218 

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 """ 

223 

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 """ 

229 

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 

236 

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. 

246 

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`. 

262 

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. 

269 

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 [] 

393 

394 

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 """ 

401 

402 task_node: TaskNode 

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

404 

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. 

408 

409 See Also 

410 -------- 

411 PipelineTaskConnections.getSpatialBoundsConnections 

412 """ 

413 

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. 

417 

418 See Also 

419 -------- 

420 PipelineTaskConnections.getTemporalBoundsConnections 

421 """ 

422 

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 """ 

427 

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 """ 

432 

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()) 

436 

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. 

440 

441 Parameters 

442 ---------- 

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

444 Data ID for this quantum. 

445 

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()) 

463 

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

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

466 a quantum. 

467 

468 Parameters 

469 ---------- 

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

471 Data ID for this quantum. 

472 

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() 

487 

488 

489class SkyPixBoundsBuilder(ABC): 

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

491 for a quantum. 

492 """ 

493 

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. 

497 

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 

506 

507 @abstractmethod 

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

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

510 

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() 

519 

520 

521class TimespanBuilder(ABC): 

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

523 for a quantum. 

524 """ 

525 

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. 

529 

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 

538 

539 @abstractmethod 

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

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

542 

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() 

550 

551 

552class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

554 needed. 

555 """ 

556 

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

558 return {} 

559 

560 

561class _TrivialTimespanBuilder(TimespanBuilder): 

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

563 

564 def finish(self) -> None: 

565 return None 

566 

567 

568class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

570 provide the only relevant spatial regions. 

571 """ 

572 

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

574 self._region = quantum_data_id.region 

575 self._dimensions = dimensions 

576 

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

578 return { 

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

580 } 

581 

582 

583class _QuantumOnlyTimespanBuilder(TimespanBuilder): 

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

585 provide the only relevant timespans. 

586 """ 

587 

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

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

590 

591 def finish(self) -> Timespan: 

592 return self._timespan 

593 

594 

595class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

597 full sky. 

598 """ 

599 

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

601 self._dimensions = dimensions 

602 

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

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

605 

606 

607class _UnboundedTimespanBuilder(TimespanBuilder): 

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

609 time. 

610 """ 

611 

612 def finish(self) -> Timespan: 

613 return Timespan(None, None) 

614 

615 

616class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

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

618 connections contribute to the spatial bounds. 

619 """ 

620 

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`. 

641 

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) 

645 

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 

654 

655 

656class _ConnectionTimespanBuilder(TimespanBuilder): 

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

658 connections contribute to the timespan. 

659 """ 

660 

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`. 

683 

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]) 

689 

690 def finish(self) -> Timespan: 

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