Coverage for python / lsst / pipe / base / connections.py: 39%

388 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-12-12 01:14 -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"""Module defining connection classes for PipelineTask.""" 

29 

30from __future__ import annotations 

31 

32__all__ = [ 

33 "AdjustQuantumHelper", 

34 "DeferredDatasetRef", 

35 "InputQuantizedConnection", 

36 "OutputQuantizedConnection", 

37 "PipelineTaskConnections", 

38 "QuantaAdjuster", 

39 "QuantizedConnection", 

40 "ScalarError", 

41 "ScalarError", 

42 "iterConnections", 

43] 

44 

45import dataclasses 

46import itertools 

47import string 

48import uuid 

49import warnings 

50from collections import UserDict, defaultdict 

51from collections.abc import Collection, Generator, Iterable, Iterator, Mapping, Sequence, Set 

52from dataclasses import dataclass 

53from types import MappingProxyType, SimpleNamespace 

54from typing import TYPE_CHECKING, Any 

55 

56from lsst.daf.butler import ( 

57 Butler, 

58 DataCoordinate, 

59 DatasetRef, 

60 DatasetType, 

61 NamedKeyDict, 

62 NamedKeyMapping, 

63 Quantum, 

64) 

65 

66from ._status import NoWorkFound 

67from .connectionTypes import BaseConnection, BaseInput, Output, PrerequisiteInput 

68 

69if TYPE_CHECKING: 

70 from .config import PipelineTaskConfig 

71 from .pipeline_graph import PipelineGraph, TaskNode 

72 from .quantum_graph_skeleton import QuantumGraphSkeleton 

73 

74 

75class ScalarError(TypeError): 

76 """Exception raised when dataset type is configured as scalar 

77 but there are multiple data IDs in a Quantum for that dataset. 

78 """ 

79 

80 

81class PipelineTaskConnectionDict(UserDict): 

82 """A special dict class used by `PipelineTaskConnectionMetaclass`. 

83 

84 This dict is used in `PipelineTaskConnection` class creation, as the 

85 dictionary that is initially used as ``__dict__``. It exists to 

86 intercept connection fields declared in a `PipelineTaskConnection`, and 

87 what name is used to identify them. The names are then added to class 

88 level list according to the connection type of the class attribute. The 

89 names are also used as keys in a class level dictionary associated with 

90 the corresponding class attribute. This information is a duplicate of 

91 what exists in ``__dict__``, but provides a simple place to lookup and 

92 iterate on only these variables. 

93 

94 Parameters 

95 ---------- 

96 *args : `~typing.Any` 

97 Passed to `dict` constructor. 

98 **kwargs : `~typing.Any` 

99 Passed to `dict` constructor. 

100 """ 

101 

102 def __init__(self, *args: Any, **kwargs: Any): 

103 super().__init__(*args, **kwargs) 

104 # Initialize class level variables used to track any declared 

105 # class level variables that are instances of 

106 # connectionTypes.BaseConnection 

107 self.data["inputs"] = set() 

108 self.data["prerequisiteInputs"] = set() 

109 self.data["outputs"] = set() 

110 self.data["initInputs"] = set() 

111 self.data["initOutputs"] = set() 

112 self.data["allConnections"] = {} 

113 

114 def __setitem__(self, name: str, value: Any) -> None: 

115 if isinstance(value, BaseConnection): 

116 if name in { 116 ↛ 126line 116 didn't jump to line 126 because the condition on line 116 was never true

117 "dimensions", 

118 "inputs", 

119 "prerequisiteInputs", 

120 "outputs", 

121 "initInputs", 

122 "initOutputs", 

123 "allConnections", 

124 }: 

125 # Guard against connections whose names are reserved. 

126 raise AttributeError(f"Connection name {name!r} is reserved for internal use.") 

127 if (previous := self.data.get(name)) is not None: 127 ↛ 130line 127 didn't jump to line 130 because the condition on line 127 was never true

128 # Guard against changing the type of an in inherited connection 

129 # by first removing it from the set it's current in. 

130 self.data[previous._connection_type_set].discard(name) 

131 object.__setattr__(value, "varName", name) 

132 self.data["allConnections"][name] = value 

133 self.data[value._connection_type_set].add(name) 

134 # defer to the default behavior 

135 super().__setitem__(name, value) 

136 

137 

138class PipelineTaskConnectionsMetaclass(type): 

139 """Metaclass used in the declaration of PipelineTaskConnections classes. 

140 

141 Parameters 

142 ---------- 

143 name : `str` 

144 Name of connection. 

145 bases : `~collections.abc.Collection` 

146 Base classes. 

147 dct : `~collections.abc.Mapping` 

148 Connections dict. 

149 **kwargs : `~typing.Any` 

150 Additional parameters. 

151 """ 

152 

153 # We can annotate these attributes as `collections.abc.Set` to discourage 

154 # undesirable modifications in type-checked code, since the internal code 

155 # modifying them is in `PipelineTaskConnectionDict` and that doesn't see 

156 # these annotations anyway. 

157 

158 dimensions: Set[str] 

159 """Set of dimension names that define the unit of work for this task. 

160 

161 Required and implied dependencies will automatically be expanded later and 

162 need not be provided. 

163 

164 This is shadowed by an instance-level attribute on 

165 `PipelineTaskConnections` instances. 

166 """ 

167 

168 inputs: Set[str] 

169 """Set with the names of all `~connectionTypes.Input` connection 

170 attributes. 

171 

172 This is updated automatically as class attributes are added. Note that 

173 this attribute is shadowed by an instance-level attribute on 

174 `PipelineTaskConnections` instances. 

175 """ 

176 

177 prerequisiteInputs: Set[str] 

178 """Set with the names of all `~connectionTypes.PrerequisiteInput` 

179 connection attributes. 

180 

181 See `inputs` for additional information. 

182 """ 

183 

184 outputs: Set[str] 

185 """Set with the names of all `~connectionTypes.Output` connection 

186 attributes. 

187 

188 See `inputs` for additional information. 

189 """ 

190 

191 initInputs: Set[str] 

192 """Set with the names of all `~connectionTypes.InitInput` connection 

193 attributes. 

194 

195 See `inputs` for additional information. 

196 """ 

197 

198 initOutputs: Set[str] 

199 """Set with the names of all `~connectionTypes.InitOutput` connection 

200 attributes. 

201 

202 See `inputs` for additional information. 

203 """ 

204 

205 allConnections: Mapping[str, BaseConnection] 

206 """Mapping containing all connection attributes. 

207 

208 See `inputs` for additional information. 

209 """ 

210 

211 def __prepare__(name, bases, **kwargs): # noqa: N804 

212 # Create an instance of our special dict to catch and track all 

213 # variables that are instances of connectionTypes.BaseConnection 

214 # Copy any existing connections from a parent class 

215 dct = PipelineTaskConnectionDict() 

216 for base in bases: 

217 if isinstance(base, PipelineTaskConnectionsMetaclass): 217 ↛ 216line 217 didn't jump to line 216 because the condition on line 217 was always true

218 for name, value in base.allConnections.items(): 218 ↛ 219line 218 didn't jump to line 219 because the loop on line 218 never started

219 dct[name] = value 

220 return dct 

221 

222 def __new__(cls, name, bases, dct, **kwargs): 

223 dimensionsValueError = TypeError( 

224 "PipelineTaskConnections class must be created with a dimensions " 

225 "attribute which is an iterable of dimension names" 

226 ) 

227 

228 if name != "PipelineTaskConnections": 

229 # Verify that dimensions are passed as a keyword in class 

230 # declaration 

231 if "dimensions" not in kwargs: 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true

232 for base in bases: 

233 if hasattr(base, "dimensions"): 

234 kwargs["dimensions"] = base.dimensions 

235 break 

236 if "dimensions" not in kwargs: 

237 raise dimensionsValueError 

238 try: 

239 if isinstance(kwargs["dimensions"], str): 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true

240 raise TypeError( 

241 "Dimensions must be iterable of dimensions, got str,possibly omitted trailing comma" 

242 ) 

243 if not isinstance(kwargs["dimensions"], Iterable): 243 ↛ 244line 243 didn't jump to line 244 because the condition on line 243 was never true

244 raise TypeError("Dimensions must be iterable of dimensions") 

245 dct["dimensions"] = set(kwargs["dimensions"]) 

246 except TypeError as exc: 

247 raise dimensionsValueError from exc 

248 # Lookup any python string templates that may have been used in the 

249 # declaration of the name field of a class connection attribute 

250 allTemplates = set() 

251 stringFormatter = string.Formatter() 

252 # Loop over all connections 

253 for obj in dct["allConnections"].values(): 

254 nameValue = obj.name 

255 # add all the parameters to the set of templates 

256 for param in stringFormatter.parse(nameValue): 

257 if param[1] is not None: 

258 allTemplates.add(param[1]) 

259 

260 # look up any template from base classes and merge them all 

261 # together 

262 mergeDict = {} 

263 mergeDeprecationsDict = {} 

264 for base in bases[::-1]: 

265 if hasattr(base, "defaultTemplates"): 

266 mergeDict.update(base.defaultTemplates) 

267 if hasattr(base, "deprecatedTemplates"): 

268 mergeDeprecationsDict.update(base.deprecatedTemplates) 

269 if "defaultTemplates" in kwargs: 

270 mergeDict.update(kwargs["defaultTemplates"]) 

271 if "deprecatedTemplates" in kwargs: 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true

272 mergeDeprecationsDict.update(kwargs["deprecatedTemplates"]) 

273 if len(mergeDict) > 0: 

274 kwargs["defaultTemplates"] = mergeDict 

275 if len(mergeDeprecationsDict) > 0: 275 ↛ 276line 275 didn't jump to line 276 because the condition on line 275 was never true

276 kwargs["deprecatedTemplates"] = mergeDeprecationsDict 

277 

278 # Verify that if templated strings were used, defaults were 

279 # supplied as an argument in the declaration of the connection 

280 # class 

281 if len(allTemplates) > 0 and "defaultTemplates" not in kwargs: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true

282 raise TypeError( 

283 "PipelineTaskConnection class contains templated attribute names, but no " 

284 "defaut templates were provided, add a dictionary attribute named " 

285 "defaultTemplates which contains the mapping between template key and value" 

286 ) 

287 if len(allTemplates) > 0: 

288 # Verify all templates have a default, and throw if they do not 

289 defaultTemplateKeys = set(kwargs["defaultTemplates"].keys()) 

290 templateDifference = allTemplates.difference(defaultTemplateKeys) 

291 if templateDifference: 291 ↛ 292line 291 didn't jump to line 292 because the condition on line 291 was never true

292 raise TypeError(f"Default template keys were not provided for {templateDifference}") 

293 # Verify that templates do not share names with variable names 

294 # used for a connection, this is needed because of how 

295 # templates are specified in an associated config class. 

296 nameTemplateIntersection = allTemplates.intersection(set(dct["allConnections"].keys())) 

297 if len(nameTemplateIntersection) > 0: 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true

298 raise TypeError( 

299 "Template parameters cannot share names with Class attributes" 

300 f" (conflicts are {nameTemplateIntersection})." 

301 ) 

302 dct["defaultTemplates"] = kwargs.get("defaultTemplates", {}) 

303 dct["deprecatedTemplates"] = kwargs.get("deprecatedTemplates", {}) 

304 

305 # Convert all the connection containers into frozensets so they cannot 

306 # be modified at the class scope 

307 for connectionName in ("inputs", "prerequisiteInputs", "outputs", "initInputs", "initOutputs"): 

308 dct[connectionName] = frozenset(dct[connectionName]) 

309 # our custom dict type must be turned into an actual dict to be used in 

310 # type.__new__ 

311 return super().__new__(cls, name, bases, dict(dct)) 

312 

313 def __init__(cls, name, bases, dct, **kwargs): 

314 # This overrides the default init to drop the kwargs argument. Python 

315 # metaclasses will have this argument set if any kwargs are passes at 

316 # class construction time, but should be consumed before calling 

317 # __init__ on the type metaclass. This is in accordance with python 

318 # documentation on metaclasses 

319 super().__init__(name, bases, dct) 

320 

321 def __call__(cls, *, config: PipelineTaskConfig | None = None) -> PipelineTaskConnections: 

322 # MyPy appears not to really understand metaclass.__call__ at all, so 

323 # we need to tell it to ignore __new__ and __init__ calls here. 

324 instance: PipelineTaskConnections = cls.__new__(cls) # type: ignore 

325 

326 # Make mutable copies of all set-like class attributes so derived 

327 # __init__ implementations can modify them in place. 

328 instance.dimensions = set(cls.dimensions) 

329 instance.inputs = set(cls.inputs) 

330 instance.prerequisiteInputs = set(cls.prerequisiteInputs) 

331 instance.outputs = set(cls.outputs) 

332 instance.initInputs = set(cls.initInputs) 

333 instance.initOutputs = set(cls.initOutputs) 

334 

335 # Set self.config. It's a bit strange that we claim to accept None but 

336 # really just raise here, but it's not worth changing now. 

337 from .config import PipelineTaskConfig # local import to avoid cycle 

338 

339 if config is None or not isinstance(config, PipelineTaskConfig): 

340 raise ValueError( 

341 "PipelineTaskConnections must be instantiated with a PipelineTaskConfig instance" 

342 ) 

343 instance.config = config 

344 

345 # Extract the template names that were defined in the config instance 

346 # by looping over the keys of the defaultTemplates dict specified at 

347 # class declaration time. 

348 templateValues = { 

349 name: getattr(config.connections, name) 

350 for name in cls.defaultTemplates # type: ignore 

351 } 

352 

353 # We now assemble a mapping of all connection instances keyed by 

354 # internal name, applying the configuration and templates to make new 

355 # configurations from the class-attribute defaults. This will be 

356 # private, but with a public read-only view. This mapping is what the 

357 # descriptor interface of the class-level attributes will return when 

358 # they are accessed on an instance. This is better than just assigning 

359 # regular instance attributes as it makes it so removed connections 

360 # cannot be accessed on instances, instead of having access to them 

361 # silent fall through to the not-removed class connection instance. 

362 instance._allConnections = {} 

363 instance.allConnections = MappingProxyType(instance._allConnections) 

364 for internal_name, connection in cls.allConnections.items(): 

365 dataset_type_name = getattr(config.connections, internal_name).format(**templateValues) 

366 instance_connection = dataclasses.replace( 

367 connection, 

368 name=dataset_type_name, 

369 doc=( 

370 connection.doc 

371 if connection.deprecated is None 

372 else f"{connection.doc}\n{connection.deprecated}" 

373 ), 

374 _deprecation_context=connection._deprecation_context, 

375 ) 

376 instance._allConnections[internal_name] = instance_connection 

377 

378 # Finally call __init__. The base class implementation does nothing; 

379 # we could have left some of the above implementation there (where it 

380 # originated), but putting it here instead makes it hard for derived 

381 # class implementors to get things into a weird state by delegating to 

382 # super().__init__ in the wrong place, or by forgetting to do that 

383 # entirely. 

384 instance.__init__(config=config) # type: ignore 

385 

386 # Derived-class implementations may have changed the contents of the 

387 # various kinds-of-connection sets; update allConnections to have keys 

388 # that are a union of all those. We get values for the new 

389 # allConnections from the attributes, since any dynamically added new 

390 # ones will not be present in the old allConnections. Typically those 

391 # getattrs will invoke the descriptors and get things from the old 

392 # allConnections anyway. After processing each set we replace it with 

393 # a frozenset. 

394 updated_all_connections = {} 

395 for attrName in ("initInputs", "prerequisiteInputs", "inputs", "initOutputs", "outputs"): 

396 updated_connection_names = getattr(instance, attrName) 

397 updated_all_connections.update( 

398 {name: getattr(instance, name) for name in updated_connection_names} 

399 ) 

400 # Setting these to frozenset is at odds with the type annotation, 

401 # but MyPy can't tell because we're using setattr, and we want to 

402 # lie to it anyway to get runtime guards against post-__init__ 

403 # mutation. 

404 setattr(instance, attrName, frozenset(updated_connection_names)) 

405 # Update the existing dict in place, since we already have a view of 

406 # that. 

407 instance._allConnections.clear() 

408 instance._allConnections.update(updated_all_connections) 

409 

410 for connection_name, obj in instance._allConnections.items(): 

411 if obj.deprecated is not None: 

412 warnings.warn( 

413 f"Connection {connection_name} with datasetType {obj.name} " 

414 f"(from {obj._deprecation_context}): {obj.deprecated}", 

415 FutureWarning, 

416 stacklevel=1, # Report from this location. 

417 ) 

418 

419 # Freeze the connection instance dimensions now. This at odds with the 

420 # type annotation, which says [mutable] `set`, just like the connection 

421 # type attributes (e.g. `inputs`, `outputs`, etc.), though MyPy can't 

422 # tell with those since we're using setattr for them. 

423 instance.dimensions = frozenset(instance.dimensions) # type: ignore 

424 

425 return instance 

426 

427 

428class QuantizedConnection(SimpleNamespace): 

429 r"""A Namespace to map defined variable names of connections to the 

430 associated `lsst.daf.butler.DatasetRef` objects. 

431 

432 This class maps the names used to define a connection on a 

433 `PipelineTaskConnections` class to the corresponding 

434 `~lsst.daf.butler.DatasetRef`\s provided by a `~lsst.daf.butler.Quantum` 

435 instance. This will be a quantum of execution based on the graph created 

436 by examining all the connections defined on the 

437 `PipelineTaskConnections` class. 

438 

439 Parameters 

440 ---------- 

441 **kwargs : `~typing.Any` 

442 Not used. 

443 """ 

444 

445 def __init__(self, **kwargs): 

446 # Create a variable to track what attributes are added. This is used 

447 # later when iterating over this QuantizedConnection instance 

448 object.__setattr__(self, "_attributes", set()) 

449 

450 def __setattr__(self, name: str, value: DatasetRef | list[DatasetRef]) -> None: 

451 # Capture the attribute name as it is added to this object 

452 self._attributes.add(name) 

453 super().__setattr__(name, value) 

454 

455 def __delattr__(self, name): 

456 object.__delattr__(self, name) 

457 self._attributes.remove(name) 

458 

459 def __len__(self) -> int: 

460 return len(self._attributes) 

461 

462 def __iter__( 

463 self, 

464 ) -> Generator[tuple[str, DatasetRef | list[DatasetRef]], None, None]: 

465 """Make an iterator for this `QuantizedConnection`. 

466 

467 Iterating over a `QuantizedConnection` will yield a tuple with the name 

468 of an attribute and the value associated with that name. This is 

469 similar to dict.items() but is on the namespace attributes rather than 

470 dict keys. 

471 """ 

472 yield from ((name, getattr(self, name)) for name in self._attributes) 

473 

474 def keys(self) -> Generator[str, None, None]: 

475 """Return an iterator over all the attributes added to a 

476 `QuantizedConnection` class. 

477 """ 

478 yield from self._attributes 

479 

480 

481class InputQuantizedConnection(QuantizedConnection): 

482 """Input variant of a `QuantizedConnection`.""" 

483 

484 pass 

485 

486 

487class OutputQuantizedConnection(QuantizedConnection): 

488 """Output variant of a `QuantizedConnection`.""" 

489 

490 pass 

491 

492 

493@dataclass(frozen=True) 

494class DeferredDatasetRef: 

495 """A wrapper class for `~lsst.daf.butler.DatasetRef` that indicates that a 

496 `PipelineTask` should receive a `~lsst.daf.butler.DeferredDatasetHandle` 

497 instead of an in-memory dataset. 

498 

499 Attributes 

500 ---------- 

501 datasetRef : `lsst.daf.butler.DatasetRef` 

502 The `lsst.daf.butler.DatasetRef` that will be eventually used to 

503 resolve a dataset. 

504 """ 

505 

506 datasetRef: DatasetRef 

507 

508 def __getattr__(self, name: str) -> Any: 

509 # make sure reduce is called on DeferredDatasetRef and not on 

510 # the DatasetRef 

511 if name in ("__reduce__", "datasetRef", "__deepcopy__"): 

512 object.__getattribute__(self, name) 

513 return getattr(self.datasetRef, name) 

514 

515 def __deepcopy__(self, memo: dict) -> DeferredDatasetRef: 

516 # dataset refs should be immutable deferred version should be too 

517 return self 

518 

519 def __reduce__(self) -> tuple: 

520 return (self.__class__, (self.datasetRef,)) 

521 

522 

523class PipelineTaskConnections(metaclass=PipelineTaskConnectionsMetaclass): 

524 """PipelineTaskConnections is a class used to declare desired IO when a 

525 PipelineTask is run by an activator. 

526 

527 Parameters 

528 ---------- 

529 config : `PipelineTaskConfig` 

530 A `PipelineTaskConfig` class instance whose class has been configured 

531 to use this `PipelineTaskConnections` class. 

532 

533 See Also 

534 -------- 

535 iterConnections : Iterator over selected connections. 

536 

537 Notes 

538 ----- 

539 ``PipelineTaskConnection`` classes are created by declaring class 

540 attributes of types defined in `lsst.pipe.base.connectionTypes` and are 

541 listed as follows: 

542 

543 * ``InitInput`` - Defines connections in a quantum graph which are used as 

544 inputs to the ``__init__`` function of the `PipelineTask` corresponding 

545 to this class 

546 * ``InitOuput`` - Defines connections in a quantum graph which are to be 

547 persisted using a butler at the end of the ``__init__`` function of the 

548 `PipelineTask` corresponding to this class. The variable name used to 

549 define this connection should be the same as an attribute name on the 

550 `PipelineTask` instance. E.g. if an ``InitOutput`` is declared with 

551 the name ``outputSchema`` in a ``PipelineTaskConnections`` class, then 

552 a `PipelineTask` instance should have an attribute 

553 ``self.outputSchema`` defined. Its value is what will be saved by the 

554 activator framework. 

555 * ``PrerequisiteInput`` - An input connection type that defines a 

556 `lsst.daf.butler.DatasetType` that must be present at execution time, 

557 but that will not be used during the course of creating the quantum 

558 graph to be executed. These most often are things produced outside the 

559 processing pipeline, such as reference catalogs. 

560 * ``Input`` - Input `lsst.daf.butler.DatasetType` objects that will be used 

561 in the ``run`` method of a `PipelineTask`. The name used to declare 

562 class attribute must match a function argument name in the ``run`` 

563 method of a `PipelineTask`. E.g. If the ``PipelineTaskConnections`` 

564 defines an ``Input`` with the name ``calexp``, then the corresponding 

565 signature should be ``PipelineTask.run(calexp, ...)`` 

566 * ``Output`` - A `lsst.daf.butler.DatasetType` that will be produced by an 

567 execution of a `PipelineTask`. The name used to declare the connection 

568 must correspond to an attribute of a `Struct` that is returned by a 

569 `PipelineTask` ``run`` method. E.g. if an output connection is 

570 defined with the name ``measCat``, then the corresponding 

571 ``PipelineTask.run`` method must return ``Struct(measCat=X,..)`` where 

572 X matches the ``storageClass`` type defined on the output connection. 

573 

574 Attributes of these types can also be created, replaced, or deleted on the 

575 `PipelineTaskConnections` instance in the ``__init__`` method, if more than 

576 just the name depends on the configuration. It is preferred to define them 

577 in the class when possible (even if configuration may cause the connection 

578 to be removed from the instance). 

579 

580 The process of declaring a ``PipelineTaskConnection`` class involves 

581 parameters passed in the declaration statement. 

582 

583 The first parameter is ``dimensions`` which is an iterable of strings which 

584 defines the unit of processing the run method of a corresponding 

585 `PipelineTask` will operate on. These dimensions must match dimensions that 

586 exist in the butler registry which will be used in executing the 

587 corresponding `PipelineTask`. The dimensions may be also modified in 

588 subclass ``__init__`` methods if they need to depend on configuration. 

589 

590 The second parameter is labeled ``defaultTemplates`` and is conditionally 

591 optional. The name attributes of connections can be specified as python 

592 format strings, with named format arguments. If any of the name parameters 

593 on connections defined in a `PipelineTaskConnections` class contain a 

594 template, then a default template value must be specified in the 

595 ``defaultTemplates`` argument. This is done by passing a dictionary with 

596 keys corresponding to a template identifier, and values corresponding to 

597 the value to use as a default when formatting the string. For example if 

598 ``ConnectionsClass.calexp.name = '{input}Coadd_calexp'`` then 

599 ``defaultTemplates`` = {'input': 'deep'}. 

600 

601 Once a `PipelineTaskConnections` class is created, it is used in the 

602 creation of a `PipelineTaskConfig`. This is further documented in the 

603 documentation of `PipelineTaskConfig`. For the purposes of this 

604 documentation, the relevant information is that the config class allows 

605 configuration of connection names by users when running a pipeline. 

606 

607 Instances of a `PipelineTaskConnections` class are used by the pipeline 

608 task execution framework to introspect what a corresponding `PipelineTask` 

609 will require, and what it will produce. 

610 

611 Examples 

612 -------- 

613 >>> from lsst.pipe.base import connectionTypes as cT 

614 >>> from lsst.pipe.base import PipelineTaskConnections 

615 >>> from lsst.pipe.base import PipelineTaskConfig 

616 >>> class ExampleConnections( 

617 ... PipelineTaskConnections, 

618 ... dimensions=("A", "B"), 

619 ... defaultTemplates={"foo": "Example"}, 

620 ... ): 

621 ... inputConnection = cT.Input( 

622 ... doc="Example input", 

623 ... dimensions=("A", "B"), 

624 ... storageClass=Exposure, 

625 ... name="{foo}Dataset", 

626 ... ) 

627 ... outputConnection = cT.Output( 

628 ... doc="Example output", 

629 ... dimensions=("A", "B"), 

630 ... storageClass=Exposure, 

631 ... name="{foo}output", 

632 ... ) 

633 >>> class ExampleConfig( 

634 ... PipelineTaskConfig, pipelineConnections=ExampleConnections 

635 ... ): 

636 ... pass 

637 >>> config = ExampleConfig() 

638 >>> config.connections.foo = Modified 

639 >>> config.connections.outputConnection = "TotallyDifferent" 

640 >>> connections = ExampleConnections(config=config) 

641 >>> assert connections.inputConnection.name == "ModifiedDataset" 

642 >>> assert connections.outputConnection.name == "TotallyDifferent" 

643 """ 

644 

645 # We annotate these attributes as mutable sets because that's what they are 

646 # inside derived ``__init__`` implementations and that's what matters most 

647 # After that's done, the metaclass __call__ makes them into frozensets, but 

648 # relatively little code interacts with them then, and that code knows not 

649 # to try to modify them without having to be told that by mypy. 

650 

651 dimensions: set[str] 

652 """Set of dimension names that define the unit of work for this task. 

653 

654 Required and implied dependencies will automatically be expanded later and 

655 need not be provided. 

656 

657 This may be replaced or modified in ``__init__`` to change the dimensions 

658 of the task. After ``__init__`` it will be a `frozenset` and may not be 

659 replaced. 

660 """ 

661 

662 inputs: set[str] 

663 """Set with the names of all `connectionTypes.Input` connection attributes. 

664 

665 This is updated automatically as class attributes are added, removed, or 

666 replaced in ``__init__``. Removing entries from this set will cause those 

667 connections to be removed after ``__init__`` completes, but this is 

668 supported only for backwards compatibility; new code should instead just 

669 delete the collection attributed directly. After ``__init__`` this will be 

670 a `frozenset` and may not be replaced. 

671 """ 

672 

673 prerequisiteInputs: set[str] 

674 """Set with the names of all `~connectionTypes.PrerequisiteInput` 

675 connection attributes. 

676 

677 See `inputs` for additional information. 

678 """ 

679 

680 outputs: set[str] 

681 """Set with the names of all `~connectionTypes.Output` connection 

682 attributes. 

683 

684 See `inputs` for additional information. 

685 """ 

686 

687 initInputs: set[str] 

688 """Set with the names of all `~connectionTypes.InitInput` connection 

689 attributes. 

690 

691 See `inputs` for additional information. 

692 """ 

693 

694 initOutputs: set[str] 

695 """Set with the names of all `~connectionTypes.InitOutput` connection 

696 attributes. 

697 

698 See `inputs` for additional information. 

699 """ 

700 

701 allConnections: Mapping[str, BaseConnection] 

702 """Mapping holding all connection attributes. 

703 

704 This is a read-only view that is automatically updated when connection 

705 attributes are added, removed, or replaced in ``__init__``. It is also 

706 updated after ``__init__`` completes to reflect changes in `inputs`, 

707 `prerequisiteInputs`, `outputs`, `initInputs`, and `initOutputs`. 

708 """ 

709 

710 _allConnections: dict[str, BaseConnection] 

711 

712 def __init__(self, *, config: PipelineTaskConfig | None = None): 

713 pass 

714 

715 def __setattr__(self, name: str, value: Any) -> None: 

716 if isinstance(value, BaseConnection): 

717 previous = self._allConnections.get(name) 

718 try: 

719 getattr(self, value._connection_type_set).add(name) 

720 except AttributeError: 

721 # Attempt to call add on a frozenset, which is what these sets 

722 # are after __init__ is done. 

723 raise TypeError("Connections objects are frozen after construction.") from None 

724 if previous is not None and value._connection_type_set != previous._connection_type_set: 

725 # Connection has changed type, e.g. Input to PrerequisiteInput; 

726 # update the sets accordingly. To be extra defensive about 

727 # multiple assignments we use the type of the previous instance 

728 # instead of assuming that's the same as the type of the self, 

729 # which is just the default. Use discard instead of remove so 

730 # manually removing from these sets first is never an error. 

731 getattr(self, previous._connection_type_set).discard(name) 

732 self._allConnections[name] = value 

733 if hasattr(self.__class__, name): 

734 # Don't actually set the attribute if this was a connection 

735 # declared in the class; in that case we let the descriptor 

736 # return the value we just added to allConnections. 

737 return 

738 # Actually add the attribute. 

739 super().__setattr__(name, value) 

740 

741 def __delattr__(self, name): 

742 """Descriptor delete method.""" 

743 previous = self._allConnections.get(name) 

744 if previous is not None: 

745 # Delete this connection's name from the appropriate set, which we 

746 # have to get from the previous instance instead of assuming it's 

747 # the same set that was appropriate for the class-level default. 

748 # Use discard instead of remove so manually removing from these 

749 # sets first is never an error. 

750 try: 

751 getattr(self, previous._connection_type_set).discard(name) 

752 except AttributeError: 

753 # Attempt to call discard on a frozenset, which is what these 

754 # sets are after __init__ is done. 

755 raise TypeError("Connections objects are frozen after construction.") from None 

756 del self._allConnections[name] 

757 if hasattr(self.__class__, name): 

758 # Don't actually delete the attribute if this was a connection 

759 # declared in the class; in that case we let the descriptor 

760 # see that it's no longer present in allConnections. 

761 return 

762 # Actually delete the attribute. 

763 super().__delattr__(name) 

764 

765 def buildDatasetRefs( 

766 self, quantum: Quantum 

767 ) -> tuple[InputQuantizedConnection, OutputQuantizedConnection]: 

768 """Build `QuantizedConnection` corresponding to input 

769 `~lsst.daf.butler.Quantum`. 

770 

771 Parameters 

772 ---------- 

773 quantum : `lsst.daf.butler.Quantum` 

774 Quantum object which defines the inputs and outputs for a given 

775 unit of processing. 

776 

777 Returns 

778 ------- 

779 retVal : `tuple` of (`InputQuantizedConnection`, \ 

780 `OutputQuantizedConnection`) 

781 Namespaces mapping attribute names 

782 (identifiers of connections) to butler references defined in the 

783 input `~lsst.daf.butler.Quantum`. 

784 """ 

785 inputDatasetRefs = InputQuantizedConnection() 

786 outputDatasetRefs = OutputQuantizedConnection() 

787 

788 # populate inputDatasetRefs from quantum inputs 

789 for attributeName in itertools.chain(self.inputs, self.prerequisiteInputs): 

790 # get the attribute identified by name 

791 attribute = getattr(self, attributeName) 

792 # if the dataset is marked to load deferred, wrap it in a 

793 # DeferredDatasetRef 

794 quantumInputRefs: list[DatasetRef] | list[DeferredDatasetRef] 

795 if attribute.deferLoad: 

796 quantumInputRefs = [ 

797 DeferredDatasetRef(datasetRef=ref) for ref in quantum.inputs[attribute.name] 

798 ] 

799 else: 

800 quantumInputRefs = list(quantum.inputs[attribute.name]) 

801 # Unpack arguments that are not marked multiples (list of 

802 # length one) 

803 if not attribute.multiple: 

804 if len(quantumInputRefs) > 1: 

805 raise ScalarError( 

806 "Received multiple datasets " 

807 f"{', '.join(str(r.dataId) for r in quantumInputRefs)} " 

808 f"for scalar connection {attributeName} " 

809 f"({quantumInputRefs[0].datasetType.name}) " 

810 f"of quantum for {quantum.taskName} with data ID {quantum.dataId}." 

811 ) 

812 if len(quantumInputRefs) == 0: 

813 continue 

814 setattr(inputDatasetRefs, attributeName, quantumInputRefs[0]) 

815 else: 

816 # Add to the QuantizedConnection identifier 

817 setattr(inputDatasetRefs, attributeName, quantumInputRefs) 

818 

819 # populate outputDatasetRefs from quantum outputs 

820 for attributeName in self.outputs: 

821 # get the attribute identified by name 

822 attribute = getattr(self, attributeName) 

823 value = quantum.outputs[attribute.name] 

824 # Unpack arguments that are not marked multiples (list of 

825 # length one) 

826 if not attribute.multiple: 

827 setattr(outputDatasetRefs, attributeName, value[0]) 

828 else: 

829 setattr(outputDatasetRefs, attributeName, value) 

830 

831 return inputDatasetRefs, outputDatasetRefs 

832 

833 def adjustQuantum( 

834 self, 

835 inputs: dict[str, tuple[BaseInput, Collection[DatasetRef]]], 

836 outputs: dict[str, tuple[Output, Collection[DatasetRef]]], 

837 label: str, 

838 data_id: DataCoordinate, 

839 ) -> tuple[ 

840 Mapping[str, tuple[BaseInput, Collection[DatasetRef]]], 

841 Mapping[str, tuple[Output, Collection[DatasetRef]]], 

842 ]: 

843 """Override to make adjustments to `lsst.daf.butler.DatasetRef` objects 

844 in the `lsst.daf.butler.Quantum` during the graph generation stage 

845 of the activator. 

846 

847 Parameters 

848 ---------- 

849 inputs : `dict` 

850 Dictionary whose keys are an input (regular or prerequisite) 

851 connection name and whose values are a tuple of the connection 

852 instance and a collection of associated 

853 `~lsst.daf.butler.DatasetRef` objects. 

854 The exact type of the nested collections is unspecified; it can be 

855 assumed to be multi-pass iterable and support `len` and ``in``, but 

856 it should not be mutated in place. In contrast, the outer 

857 dictionaries are guaranteed to be temporary copies that are true 

858 `dict` instances, and hence may be modified and even returned; this 

859 is especially useful for delegating to `super` (see notes below). 

860 outputs : `~collections.abc.Mapping` 

861 Mapping of output datasets, with the same structure as ``inputs``. 

862 label : `str` 

863 Label for this task in the pipeline (should be used in all 

864 diagnostic messages). 

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

866 Data ID for this quantum in the pipeline (should be used in all 

867 diagnostic messages). 

868 

869 Returns 

870 ------- 

871 adjusted_inputs : `~collections.abc.Mapping` 

872 Mapping of the same form as ``inputs`` with updated containers of 

873 input `~lsst.daf.butler.DatasetRef` objects. Connections that are 

874 not changed should not be returned at all. Datasets may only be 

875 removed, not added. Nested collections may be of any multi-pass 

876 iterable type, and the order of iteration will set the order of 

877 iteration within `PipelineTask.runQuantum`. 

878 adjusted_outputs : `~collections.abc.Mapping` 

879 Mapping of updated output datasets, with the same structure and 

880 interpretation as ``adjusted_inputs``. 

881 

882 Raises 

883 ------ 

884 ScalarError 

885 Raised if any `Input` or `PrerequisiteInput` connection has 

886 ``multiple`` set to `False`, but multiple datasets. 

887 NoWorkFound 

888 Raised to indicate that this quantum should not be run; not enough 

889 datasets were found for a regular `Input` connection, and the 

890 quantum should be pruned or skipped. 

891 FileNotFoundError 

892 Raised to cause QuantumGraph generation to fail (with the message 

893 included in this exception); not enough datasets were found for a 

894 `PrerequisiteInput` connection. 

895 

896 Notes 

897 ----- 

898 The base class implementation performs important checks. It always 

899 returns an empty mapping (i.e. makes no adjustments). It should 

900 always called be via `super` by custom implementations, ideally at the 

901 end of the custom implementation with already-adjusted mappings when 

902 any datasets are actually dropped, e.g.: 

903 

904 .. code-block:: python 

905 

906 def adjustQuantum(self, inputs, outputs, label, data_id): 

907 # Filter out some dataset refs for one connection. 

908 connection, old_refs = inputs["my_input"] 

909 new_refs = [ref for ref in old_refs if ...] 

910 adjusted_inputs = {"my_input", (connection, new_refs)} 

911 # Update the original inputs so we can pass them to super. 

912 inputs.update(adjusted_inputs) 

913 # Can ignore outputs from super because they are guaranteed 

914 # to be empty. 

915 super().adjustQuantum(inputs, outputs, label_data_id) 

916 # Return only the connections we modified. 

917 return adjusted_inputs, {} 

918 

919 Removing outputs here is guaranteed to affect what is actually 

920 passed to `PipelineTask.runQuantum`, but its effect on the larger 

921 graph may be deferred to execution, depending on the context in 

922 which `adjustQuantum` is being run: if one quantum removes an output 

923 that is needed by a second quantum as input, the second quantum may not 

924 be adjusted (and hence pruned or skipped) until that output is actually 

925 found to be missing at execution time. 

926 

927 Tasks that desire zip-iteration consistency between any combinations of 

928 connections that have the same data ID should generally implement 

929 `adjustQuantum` to achieve this, even if they could also run that 

930 logic during execution; this allows the system to see outputs that will 

931 not be produced because the corresponding input is missing as early as 

932 possible. 

933 """ 

934 for name, (input_connection, refs) in inputs.items(): 

935 dataset_type_name = input_connection.name 

936 if not input_connection.multiple and len(refs) > 1: 

937 raise ScalarError( 

938 f"Found multiple datasets {', '.join(str(r.dataId) for r in refs)} " 

939 f"for non-multiple input connection {label}.{name} ({dataset_type_name}) " 

940 f"for quantum data ID {data_id}." 

941 ) 

942 if len(refs) < input_connection.minimum: 

943 if isinstance(input_connection, PrerequisiteInput): 

944 # This branch should only be possible during QG generation, 

945 # or if someone deleted the dataset between making the QG 

946 # and trying to run it. Either one should be a hard error. 

947 raise FileNotFoundError( 

948 f"Not enough datasets ({len(refs)}) found for non-optional connection {label}.{name} " 

949 f"({dataset_type_name}) with minimum={input_connection.minimum} for quantum data ID " 

950 f"{data_id}." 

951 ) 

952 else: 

953 raise NoWorkFound(label, name, input_connection) 

954 for name, (output_connection, refs) in outputs.items(): 

955 dataset_type_name = output_connection.name 

956 if not output_connection.multiple and len(refs) > 1: 

957 raise ScalarError( 

958 f"Found multiple datasets {', '.join(str(r.dataId) for r in refs)} " 

959 f"for non-multiple output connection {label}.{name} ({dataset_type_name}) " 

960 f"for quantum data ID {data_id}." 

961 ) 

962 return {}, {} 

963 

964 def getSpatialBoundsConnections(self) -> Iterable[str]: 

965 """Return the names of regular input and output connections whose data 

966 IDs should be used to compute the spatial bounds of this task's quanta. 

967 

968 The spatial bound for a quantum is defined as the union of the regions 

969 of all data IDs of all connections returned here, along with the region 

970 of the quantum data ID (if the task has spatial dimensions). 

971 

972 Returns 

973 ------- 

974 connection_names : `collections.abc.Iterable` [ `str` ] 

975 Names of collections with spatial dimensions. These are the 

976 task-internal connection names, not butler dataset type names. 

977 

978 Notes 

979 ----- 

980 The spatial bound is used to search for prerequisite inputs that have 

981 skypix dimensions. The default implementation returns an empty 

982 iterable, which is usually sufficient for tasks with spatial 

983 dimensions, but if a task's inputs or outputs are associated with 

984 spatial regions that extend beyond the quantum data ID's region, this 

985 method may need to be overridden to expand the set of prerequisite 

986 inputs found. 

987 

988 Tasks that do not have spatial dimensions that have skypix prerequisite 

989 inputs should always override this method, as the default spatial 

990 bounds otherwise cover the full sky. 

991 """ 

992 return () 

993 

994 def getTemporalBoundsConnections(self) -> Iterable[str]: 

995 """Return the names of regular input and output connections whose data 

996 IDs should be used to compute the temporal bounds of this task's 

997 quanta. 

998 

999 The temporal bound for a quantum is defined as the union of the 

1000 timespans of all data IDs of all connections returned here, along with 

1001 the timespan of the quantum data ID (if the task has temporal 

1002 dimensions). 

1003 

1004 Returns 

1005 ------- 

1006 connection_names : `collections.abc.Iterable` [ `str` ] 

1007 Names of collections with temporal dimensions. These are the 

1008 task-internal connection names, not butler dataset type names. 

1009 

1010 Notes 

1011 ----- 

1012 The temporal bound is used to search for prerequisite inputs that are 

1013 calibration datasets. The default implementation returns an empty 

1014 iterable, which is usually sufficient for tasks with temporal 

1015 dimensions, but if a task's inputs or outputs are associated with 

1016 timespans that extend beyond the quantum data ID's timespan, this 

1017 method may need to be overridden to expand the set of prerequisite 

1018 inputs found. 

1019 

1020 Tasks that do not have temporal dimensions that do not implement this 

1021 method will use an infinite timespan for any calibration lookups. 

1022 """ 

1023 return () 

1024 

1025 def adjust_all_quanta(self, adjuster: QuantaAdjuster) -> None: 

1026 """Customize the set of quanta predicted for this task during quantum 

1027 graph generation. 

1028 

1029 Parameters 

1030 ---------- 

1031 adjuster : `QuantaAdjuster` 

1032 A helper object that implementations can use to modify the 

1033 under-construction quantum graph. 

1034 

1035 Notes 

1036 ----- 

1037 This hook is called before `adjustQuantum`, which is where built-in 

1038 checks for `NoWorkFound` cases and missing prerequisites are handled. 

1039 This means that the set of preliminary quanta seen by this method could 

1040 include some that would normally be dropped later. 

1041 """ 

1042 pass 

1043 

1044 

1045def iterConnections( 

1046 connections: PipelineTaskConnections, connectionType: str | Iterable[str] 

1047) -> Generator[BaseConnection, None, None]: 

1048 """Create an iterator over the selected connections type which yields 

1049 all the defined connections of that type. 

1050 

1051 Parameters 

1052 ---------- 

1053 connections : `PipelineTaskConnections` 

1054 An instance of a `PipelineTaskConnections` object that will be iterated 

1055 over. 

1056 connectionType : `str` 

1057 The type of connections to iterate over, valid values are inputs, 

1058 outputs, prerequisiteInputs, initInputs, initOutputs. 

1059 

1060 Yields 

1061 ------ 

1062 connection: `~.connectionTypes.BaseConnection` 

1063 A connection defined on the input connections object of the type 

1064 supplied. The yielded value Will be an derived type of 

1065 `~.connectionTypes.BaseConnection`. 

1066 """ 

1067 if isinstance(connectionType, str): 

1068 connectionType = (connectionType,) 

1069 for name in itertools.chain.from_iterable(getattr(connections, ct) for ct in connectionType): 

1070 yield getattr(connections, name) 

1071 

1072 

1073@dataclass 

1074class AdjustQuantumHelper: 

1075 """Helper class for calling `PipelineTaskConnections.adjustQuantum`. 

1076 

1077 This class holds `inputs` and `outputs` mappings in the form used by 

1078 `lsst.daf.butler.Quantum` and execution harness code, i.e. with 

1079 `~lsst.daf.butler.DatasetType` keys, translating them to and from the 

1080 connection-oriented mappings used inside `PipelineTaskConnections`. 

1081 """ 

1082 

1083 inputs: NamedKeyMapping[DatasetType, Sequence[DatasetRef]] 

1084 """Mapping of regular input and prerequisite input datasets, grouped by 

1085 `~lsst.daf.butler.DatasetType`. 

1086 """ 

1087 

1088 outputs: NamedKeyMapping[DatasetType, Sequence[DatasetRef]] 

1089 """Mapping of output datasets, grouped by `~lsst.daf.butler.DatasetType`. 

1090 """ 

1091 

1092 inputs_adjusted: bool = False 

1093 """Whether any inputs were removed in the last call to `adjust_in_place`. 

1094 """ 

1095 

1096 outputs_adjusted: bool = False 

1097 """Whether any outputs were removed in the last call to `adjust_in_place`. 

1098 """ 

1099 

1100 def adjust_in_place( 

1101 self, 

1102 connections: PipelineTaskConnections, 

1103 label: str, 

1104 data_id: DataCoordinate, 

1105 ) -> None: 

1106 """Call `~PipelineTaskConnections.adjustQuantum` and update ``self`` 

1107 with its results. 

1108 

1109 Parameters 

1110 ---------- 

1111 connections : `PipelineTaskConnections` 

1112 Instance on which to call `~PipelineTaskConnections.adjustQuantum`. 

1113 label : `str` 

1114 Label for this task in the pipeline (should be used in all 

1115 diagnostic messages). 

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

1117 Data ID for this quantum in the pipeline (should be used in all 

1118 diagnostic messages). 

1119 """ 

1120 # Translate self's DatasetType-keyed, Quantum-oriented mappings into 

1121 # connection-keyed, PipelineTask-oriented mappings. 

1122 inputs_by_connection: dict[str, tuple[BaseInput, tuple[DatasetRef, ...]]] = {} 

1123 outputs_by_connection: dict[str, tuple[Output, tuple[DatasetRef, ...]]] = {} 

1124 for name in itertools.chain(connections.inputs, connections.prerequisiteInputs): 

1125 connection = getattr(connections, name) 

1126 dataset_type_name = connection.name 

1127 inputs_by_connection[name] = (connection, tuple(self.inputs.get(dataset_type_name, ()))) 

1128 for name in itertools.chain(connections.outputs): 

1129 connection = getattr(connections, name) 

1130 dataset_type_name = connection.name 

1131 outputs_by_connection[name] = (connection, tuple(self.outputs.get(dataset_type_name, ()))) 

1132 # Actually call adjustQuantum. 

1133 # MyPy correctly complains that this call is not quite legal, but the 

1134 # method docs explain exactly what's expected and it's the behavior we 

1135 # want. It'd be nice to avoid this if we ever have to change the 

1136 # interface anyway, but not an immediate problem. 

1137 adjusted_inputs_by_connection, adjusted_outputs_by_connection = connections.adjustQuantum( 

1138 inputs_by_connection, # type: ignore 

1139 outputs_by_connection, # type: ignore 

1140 label, 

1141 data_id, 

1142 ) 

1143 # Translate adjustments to DatasetType-keyed, Quantum-oriented form, 

1144 # installing new mappings in self if necessary. 

1145 if adjusted_inputs_by_connection: 

1146 adjusted_inputs = NamedKeyDict[DatasetType, tuple[DatasetRef, ...]](self.inputs) 

1147 for name, (connection, updated_refs) in adjusted_inputs_by_connection.items(): 

1148 dataset_type_name = connection.name 

1149 if not set(updated_refs).issubset(self.inputs[dataset_type_name]): 

1150 raise RuntimeError( 

1151 f"adjustQuantum implementation for task with label {label} returned {name} " 

1152 f"({dataset_type_name}) input datasets that are not a subset of those " 

1153 f"it was given for data ID {data_id}." 

1154 ) 

1155 adjusted_inputs[dataset_type_name] = tuple(updated_refs) 

1156 self.inputs = adjusted_inputs.freeze() 

1157 self.inputs_adjusted = True 

1158 else: 

1159 self.inputs_adjusted = False 

1160 if adjusted_outputs_by_connection: 

1161 adjusted_outputs = NamedKeyDict[DatasetType, tuple[DatasetRef, ...]](self.outputs) 

1162 for name, (connection, updated_refs) in adjusted_outputs_by_connection.items(): 

1163 dataset_type_name = connection.name 

1164 if not set(updated_refs).issubset(self.outputs[dataset_type_name]): 

1165 raise RuntimeError( 

1166 f"adjustQuantum implementation for task with label {label} returned {name} " 

1167 f"({dataset_type_name}) output datasets that are not a subset of those " 

1168 f"it was given for data ID {data_id}." 

1169 ) 

1170 adjusted_outputs[dataset_type_name] = tuple(updated_refs) 

1171 self.outputs = adjusted_outputs.freeze() 

1172 self.outputs_adjusted = True 

1173 else: 

1174 self.outputs_adjusted = False 

1175 

1176 

1177class QuantaAdjuster: 

1178 """A helper class for the `PipelineTaskConnections.adjust_all_quanta` hook. 

1179 

1180 Parameters 

1181 ---------- 

1182 task_label : `str` 

1183 Label of the task whose quanta are being adjusted. 

1184 pipeline_graph : `pipeline_graph.PipelineGraph` 

1185 Pipeline graph the quantum graph is being built from. 

1186 skeleton : `quantum_graph_skeleton.QuantumGraphSkeleton` 

1187 Under-construction quantum graph that will be modified in place. 

1188 butler : `lsst.daf.butler.Butler` 

1189 Read-only instance with its default collection search path set to the 

1190 input collections passed to the quantum-graph builder. 

1191 """ 

1192 

1193 def __init__( 

1194 self, task_label: str, pipeline_graph: PipelineGraph, skeleton: QuantumGraphSkeleton, butler: Butler 

1195 ): 

1196 self._task_node = pipeline_graph.tasks[task_label] 

1197 self._pipeline_graph = pipeline_graph 

1198 self._skeleton = skeleton 

1199 self._n_removed = 0 

1200 self._butler = butler 

1201 

1202 @property 

1203 def task_label(self) -> str: 

1204 """The label this task has been configured with.""" 

1205 return self._task_node.label 

1206 

1207 @property 

1208 def task_node(self) -> TaskNode: 

1209 """The node for this task in the pipeline graph.""" 

1210 return self._task_node 

1211 

1212 @property 

1213 def butler(self) -> Butler: 

1214 """Read-only instance with its default collection search path set to 

1215 the input collections passed to the quantum-graph builder. 

1216 """ 

1217 return self._butler 

1218 

1219 def iter_data_ids(self) -> Iterator[DataCoordinate]: 

1220 """Iterate over the data IDs of all quanta for this task. 

1221 

1222 Returns 

1223 ------- 

1224 data_ids : `~collections.abc.Iterator` [ \ 

1225 `~lsst.daf.butler.DataCoordinate` ] 

1226 Data IDs. These are minimal data IDs without dimension records or 

1227 implied values; use `expand_quantum_data_id` to get a full data ID 

1228 when needed. 

1229 """ 

1230 for key in self._skeleton.get_quanta(self._task_node.label): 

1231 yield DataCoordinate.from_required_values(self._task_node.dimensions, key.data_id_values) 

1232 

1233 def remove_quantum(self, data_id: DataCoordinate) -> None: 

1234 """Remove a quantum from the graph. 

1235 

1236 Parameters 

1237 ---------- 

1238 data_id : `~lsst.daf.butler.DataCoordinate` 

1239 Data ID of the quantum to remove. All outputs will be removed as 

1240 well. 

1241 """ 

1242 from .quantum_graph_skeleton import QuantumKey 

1243 

1244 self._skeleton.remove_quantum_node( 

1245 QuantumKey(self._task_node.label, data_id.required_values), remove_outputs=True 

1246 ) 

1247 self._n_removed += 1 

1248 

1249 def get_inputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]: 

1250 """Return the data IDs of all regular inputs to a quantum. 

1251 

1252 Parameters 

1253 ---------- 

1254 quantum_data_id : `~lsst.daf.butler.DataCoordinate` 

1255 Data ID of the quantum to get the inputs of. 

1256 

1257 Returns 

1258 ------- 

1259 inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ] 

1260 Data IDs of inputs, keyed by the connection name (the internal task 

1261 name, not the dataset type name). This only contains regular 

1262 inputs, not init-inputs or prerequisite inputs. 

1263 

1264 Notes 

1265 ----- 

1266 If two connections have the same dataset type, the current 

1267 implementation assumes the set of datasets is the same for the two 

1268 connections. This limitation may be removed in the future. 

1269 """ 

1270 from .quantum_graph_skeleton import DatasetKey, QuantumKey 

1271 

1272 by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list) 

1273 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values) 

1274 for dataset_key in self._skeleton.iter_inputs_of(quantum_key): 

1275 if not isinstance(dataset_key, DatasetKey): 

1276 continue 

1277 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name] 

1278 by_dataset_type_name[dataset_key.parent_dataset_type_name].append( 

1279 DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values) 

1280 ) 

1281 return { 

1282 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name] 

1283 for edge in self._task_node.iter_all_inputs() 

1284 } 

1285 

1286 def get_prerequisite_inputs( 

1287 self, 

1288 quantum_data_id: DataCoordinate, 

1289 ) -> dict[str, dict[uuid.UUID, DataCoordinate]]: 

1290 """Return the data IDs of all prerequisite inputs to a quantum. 

1291 

1292 Parameters 

1293 ---------- 

1294 quantum_data_id : `~lsst.daf.butler.DataCoordinate` 

1295 Data ID of the quantum to get the inputs of. 

1296 

1297 Returns 

1298 ------- 

1299 inputs : `dict` [ `str`, \ 

1300 `dict` [ `uuid.UUID`, `~lsst.daf.butler.DataCoordinate` ] ] 

1301 Dataset IDs and and data IDs of prerequisite inputs, keyed by the 

1302 connection name (the internal task name, not the dataset type 

1303 name). This only contains prerequisite inputs, not init-inputs or 

1304 regular inputs. 

1305 

1306 Notes 

1307 ----- 

1308 If two connections have the same dataset type, the current 

1309 implementation assumes the set of datasets is the same for the two 

1310 connections. This limitation may be removed in the future. 

1311 

1312 Unlike regular inputs, prerequisite inputs are not looked up from input 

1313 collections or indexed by data ID. Instead, they are uniquely 

1314 identified by dataset UUID and reused directly between quanta. 

1315 """ 

1316 from .quantum_graph_skeleton import PrerequisiteDatasetKey, QuantumKey 

1317 

1318 by_dataset_type_name: defaultdict[str, dict[uuid.UUID, DataCoordinate]] = defaultdict(dict) 

1319 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values) 

1320 

1321 for dataset_key in self._skeleton.iter_inputs_of(quantum_key): 

1322 if not isinstance(dataset_key, PrerequisiteDatasetKey): 

1323 continue 

1324 by_dataset_type_name[dataset_key.parent_dataset_type_name][ 

1325 uuid.UUID(bytes=dataset_key.dataset_id_bytes) 

1326 ] = self._skeleton.get_data_id(dataset_key) 

1327 return { 

1328 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name] 

1329 for edge in self._task_node.iter_all_inputs() 

1330 } 

1331 

1332 def get_outputs(self, quantum_data_id: DataCoordinate) -> dict[str, list[DataCoordinate]]: 

1333 """Return the data IDs of all regular outputs to a quantum. 

1334 

1335 Parameters 

1336 ---------- 

1337 quantum_data_id : `~lsst.daf.butler.DataCoordinate` 

1338 Data ID of the quantum to get the outputs of. 

1339 

1340 Returns 

1341 ------- 

1342 inputs : `dict` [ `str`, `list` [ `~lsst.daf.butler.DataCoordinate` ] ] 

1343 Data IDs of inputs, keyed by the connection name (the internal task 

1344 name, not the dataset type name). This only contains regular 

1345 outputs, not init-outputs or log or metadata outputs. 

1346 

1347 Notes 

1348 ----- 

1349 If two connections have the same dataset type, the current 

1350 implementation assumes the set of datasets is the same for the two 

1351 connections. This limitation may be removed in the future. 

1352 """ 

1353 from .quantum_graph_skeleton import QuantumKey 

1354 

1355 by_dataset_type_name: defaultdict[str, list[DataCoordinate]] = defaultdict(list) 

1356 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values) 

1357 for dataset_key in self._skeleton.iter_outputs_of(quantum_key): 

1358 dataset_type_node = self._pipeline_graph.dataset_types[dataset_key.parent_dataset_type_name] 

1359 by_dataset_type_name[dataset_key.parent_dataset_type_name].append( 

1360 DataCoordinate.from_required_values(dataset_type_node.dimensions, dataset_key.data_id_values) 

1361 ) 

1362 return { 

1363 edge.connection_name: by_dataset_type_name[edge.parent_dataset_type_name] 

1364 for edge in self._task_node.outputs.values() 

1365 } 

1366 

1367 def add_input( 

1368 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate 

1369 ) -> None: 

1370 """Add a new input to a quantum. 

1371 

1372 Parameters 

1373 ---------- 

1374 quantum_data_id : `~lsst.daf.butler.DataCoordinate` 

1375 Data ID of the quantum to add an input to. 

1376 connection_name : `str` 

1377 Name of the connection (the task-internal name, not the butler 

1378 dataset type name). 

1379 dataset_data_id : `~lsst.daf.butler.DataCoordinate` 

1380 Data ID of the input dataset. Must already exist in the graph 

1381 as an input to a different quantum of this task, and must be a 

1382 regular input, not a prerequisite input or init-input. 

1383 

1384 Notes 

1385 ----- 

1386 If two connections have the same dataset type, the current 

1387 implementation assumes the set of datasets is the same for the two 

1388 connections. This limitation may be removed in the future. 

1389 """ 

1390 from .quantum_graph_skeleton import DatasetKey, QuantumKey 

1391 

1392 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values) 

1393 read_edge = self._task_node.inputs[connection_name] 

1394 dataset_key = DatasetKey(read_edge.parent_dataset_type_name, dataset_data_id.required_values) 

1395 if dataset_key not in self._skeleton: 

1396 raise LookupError( 

1397 f"Dataset {read_edge.parent_dataset_type_name}@{dataset_data_id} is not already in the graph." 

1398 ) 

1399 self._skeleton.add_input_edge(quantum_key, dataset_key) 

1400 

1401 def add_prerequisite_input( 

1402 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_uuid: uuid.UUID 

1403 ) -> None: 

1404 """Add a new prerequisite input to a quantum. 

1405 

1406 Parameters 

1407 ---------- 

1408 quantum_data_id : `~lsst.daf.butler.DataCoordinate` 

1409 Data ID of the quantum to add an input to. 

1410 connection_name : `str` 

1411 Name of the connection (the task-internal name, not the butler 

1412 dataset type name). 

1413 dataset_uuid : `uuid.UUID` 

1414 UUID of the prerequisite input dataset. Must already exist in the 

1415 graph as an input to a different quantum of this task, and must be 

1416 a prerequisite input, not a regular input or init-input. 

1417 

1418 Notes 

1419 ----- 

1420 If two connections have the same dataset type, the current 

1421 implementation assumes the set of datasets is the same for the two 

1422 connections. This limitation may be removed in the future. 

1423 

1424 Unlike regular inputs, prerequisite inputs are not looked up from input 

1425 collections or indexed by data ID. Instead, they are uniquely 

1426 identified by dataset UUID and reused directly between quanta. 

1427 """ 

1428 from .quantum_graph_skeleton import PrerequisiteDatasetKey, QuantumKey 

1429 

1430 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values) 

1431 read_edge = self._task_node.prerequisite_inputs[connection_name] 

1432 dataset_key = PrerequisiteDatasetKey(read_edge.parent_dataset_type_name, dataset_uuid.bytes) 

1433 if dataset_key not in self._skeleton: 

1434 raise LookupError( 

1435 f"Prerequisite Dataset {read_edge.parent_dataset_type_name}@{dataset_uuid} " 

1436 "is not already in the graph." 

1437 ) 

1438 self._skeleton.add_input_edge(quantum_key, dataset_key) 

1439 

1440 def move_output( 

1441 self, quantum_data_id: DataCoordinate, connection_name: str, dataset_data_id: DataCoordinate 

1442 ) -> None: 

1443 """Remove an output of one quantum and make it a new output of another. 

1444 

1445 Parameters 

1446 ---------- 

1447 quantum_data_id : `~lsst.daf.butler.DataCoordinate` 

1448 Data ID of the quantum to move the output to. 

1449 connection_name : `str` 

1450 Name of the connection (the task-internal name, not the butler 

1451 dataset type name). 

1452 dataset_data_id : `~lsst.daf.butler.DataCoordinate` 

1453 Data ID of the output dataset. Must already exist in the graph 

1454 as an output of a different quantum of this task. 

1455 """ 

1456 from .quantum_graph_skeleton import DatasetKey, QuantumKey 

1457 

1458 quantum_key = QuantumKey(self._task_node.label, quantum_data_id.required_values) 

1459 write_edge = self._task_node.outputs[connection_name] 

1460 dataset_key = DatasetKey(write_edge.parent_dataset_type_name, dataset_data_id.required_values) 

1461 if dataset_key not in self._skeleton: 

1462 raise LookupError( 

1463 f"Dataset {write_edge.parent_dataset_type_name}@{dataset_data_id} is " 

1464 "not already in the graph." 

1465 ) 

1466 self._skeleton.remove_output_edge(dataset_key) 

1467 self._skeleton.add_output_edge(quantum_key, dataset_key) 

1468 

1469 def expand_quantum_data_id(self, data_id: DataCoordinate) -> DataCoordinate: 

1470 """Expand a quantum data ID to include implied values and dimension 

1471 records. 

1472 

1473 Parameters 

1474 ---------- 

1475 data_id : `~lsst.daf.butler.DataCoordinate` 

1476 A data ID of a quantum already in the graph. 

1477 

1478 Returns 

1479 ------- 

1480 expanded_data_id : `~lsst.daf.butler.DataCoordinate` 

1481 The same data ID, with implied values included and dimension 

1482 records attached. 

1483 """ 

1484 from .quantum_graph_skeleton import QuantumKey 

1485 

1486 return self._skeleton.get_data_id(QuantumKey(self._task_node.label, data_id.required_values)) 

1487 

1488 @property 

1489 def n_removed(self) -> int: 

1490 """The number of quanta that have been removed by this helper.""" 

1491 return self._n_removed