Coverage for python / lsst / pipe / base / pipeline_graph / _edges.py: 28%
211 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/>.
27from __future__ import annotations
29__all__ = ("Edge", "ReadEdge", "WriteEdge")
31from abc import ABC, abstractmethod
32from collections.abc import Callable, Mapping, Sequence
33from typing import Any, ClassVar, Self, TypeVar
35from lsst.daf.butler import DatasetRef, DatasetType, DimensionUniverse, StorageClassFactory
36from lsst.daf.butler.registry import MissingDatasetTypeError
37from lsst.utils.classes import immutable
39from ..connectionTypes import BaseConnection
40from ._exceptions import ConnectionTypeConsistencyError, IncompatibleDatasetTypeError
41from ._nodes import NodeKey, NodeType
43_S = TypeVar("_S", bound="Edge")
46@immutable
47class Edge(ABC):
48 """Base class for edges in a pipeline graph.
50 This represents the link between a task node and an input or output dataset
51 type.
53 Parameters
54 ----------
55 task_key : `NodeKey`
56 Key for the task node this edge is connected to.
57 dataset_type_key : `NodeKey`
58 Key for the dataset type node this edge is connected to.
59 storage_class_name : `str`
60 Name of the dataset type's storage class as seen by the task.
61 connection_name : `str`
62 Internal name for the connection as seen by the task.
63 is_calibration : `bool`
64 Whether this dataset type can be included in
65 `~lsst.daf.butler.CollectionType.CALIBRATION` collections.
66 raw_dimensions : `frozenset` [ `str` ]
67 Raw dimensions from the connection definition.
68 """
70 def __init__(
71 self,
72 *,
73 task_key: NodeKey,
74 dataset_type_key: NodeKey,
75 storage_class_name: str,
76 connection_name: str,
77 is_calibration: bool,
78 raw_dimensions: frozenset[str],
79 ):
80 self.task_key = task_key
81 self.dataset_type_key = dataset_type_key
82 self.connection_name = connection_name
83 self.storage_class_name = storage_class_name
84 self.is_calibration = is_calibration
85 self.raw_dimensions = raw_dimensions
87 INIT_TO_TASK_NAME: ClassVar[str] = "INIT"
88 """Edge key for the special edge that connects a task init node to the
89 task node itself (for regular edges, this would be the connection name).
90 """
92 task_key: NodeKey
93 """Task part of the key for this edge in networkx graphs (`NodeKey`)."""
95 dataset_type_key: NodeKey
96 """Task part of the key for this edge in networkx graphs (`NodeKey`)."""
98 connection_name: str
99 """Name used by the task to refer to this dataset type (`str`)."""
101 storage_class_name: str
102 """Storage class expected by this task (`str`).
104 If `ReadEdge.component` is not `None`, this is the component storage class,
105 not the parent storage class.
106 """
108 is_calibration: bool
109 """Whether this dataset type can be included in
110 `~lsst.daf.butler.CollectionType.CALIBRATION` collections.
111 """
113 raw_dimensions: frozenset[str]
114 """Raw dimensions in the task declaration (`frozenset` [`str`]).
116 This can only be used safely for partial comparisons: two edges with the
117 same ``raw_dimensions`` (and the same parent dataset type name) always have
118 the same resolved dimensions, but edges with different ``raw_dimensions``
119 may also have the same resolvd dimensions.
120 """
122 @property
123 def is_init(self) -> bool:
124 """Whether this dataset is read or written when the task is
125 constructed, not when it is run.
126 """
127 return self.task_key.node_type is NodeType.TASK_INIT
129 @property
130 def task_label(self) -> str:
131 """Label of the task (`str`)."""
132 return str(self.task_key)
134 @property
135 def parent_dataset_type_name(self) -> str:
136 """Name of the parent dataset type (`str`).
138 All dataset type nodes in a pipeline graph are for parent dataset
139 types; components are represented by additional `ReadEdge` state.
140 """
141 return str(self.dataset_type_key)
143 @property
144 @abstractmethod
145 def nodes(self) -> tuple[NodeKey, NodeKey]:
146 """The directed pair of `NodeKey` instances this edge connects.
148 This tuple is ordered in the same direction as the pipeline flow:
149 `task_key` precedes `dataset_type_key` for writes, and the
150 reverse is true for reads.
151 """
152 raise NotImplementedError()
154 @property
155 def key(self) -> tuple[NodeKey, NodeKey, str]:
156 """Ordered tuple of node keys and connection name that uniquely
157 identifies this edge in a pipeline graph (`NodeKey`, `NodeKey`, `str`).
159 The nodes are ordered in the same sense as for `nodes`.
160 """
161 return self.nodes + (self.connection_name,)
163 def __repr__(self) -> str:
164 return f"{self.nodes[0]} -> {self.nodes[1]} ({self.connection_name})"
166 @property
167 def dataset_type_name(self) -> str:
168 """Dataset type name seen by the task (`str`).
170 This defaults to the parent dataset type name, which is appropriate
171 for all writes and most reads.
172 """
173 return self.parent_dataset_type_name
175 def diff(self: _S, other: _S, connection_type: str = "connection") -> list[str]:
176 """Compare this edge to another one from a possibly-different
177 configuration of the same task label.
179 Parameters
180 ----------
181 other : `Edge`
182 Another edge of the same type to compare to.
183 connection_type : `str`
184 Human-readable name of the connection type of this edge (e.g.
185 "init input", "output") for use in returned messages.
187 Returns
188 -------
189 differences : `list` [ `str` ]
190 List of string messages describing differences between ``self`` and
191 ``other``. Will be empty if ``self == other`` or if the only
192 difference is in the task label or connection name (which are not
193 checked). Messages will use 'A' to refer to ``self`` and 'B' to
194 refer to ``other``.
195 """
196 result = []
197 if self.dataset_type_name != other.dataset_type_name:
198 result.append(
199 f"{connection_type.capitalize()} {self.task_label}.{self.connection_name} has dataset type "
200 f"{self.dataset_type_name!r} in A, but {other.dataset_type_name!r} in B."
201 )
202 if self.storage_class_name != other.storage_class_name:
203 result.append(
204 f"{connection_type.capitalize()} {self.task_label}.{self.connection_name} has storage class "
205 f"{self.storage_class_name!r} in A, but {other.storage_class_name!r} in B."
206 )
207 if self.raw_dimensions != other.raw_dimensions:
208 result.append(
209 f"{connection_type.capitalize()} {self.task_label}.{self.connection_name} has raw dimensions "
210 f"{set(self.raw_dimensions)} in A, but {set(other.raw_dimensions)} in B "
211 "(differences in raw dimensions may not lead to differences in resolved dimensions, "
212 "but this cannot be checked without re-resolving the dataset type)."
213 )
214 if self.is_calibration != other.is_calibration:
215 result.append(
216 f"{connection_type.capitalize()} {self.task_label}.{self.connection_name} is marked as a "
217 f"calibration {'in A but not in B' if self.is_calibration else 'in B but not in A'}."
218 )
219 return result
221 @abstractmethod
222 def adapt_dataset_type(self, dataset_type: DatasetType) -> DatasetType:
223 """Transform the graph's definition of a dataset type (parent, with the
224 registry or producer's storage class) to the one seen by this task.
226 Parameters
227 ----------
228 dataset_type : `~lsst.daf.butler.DatasetType`
229 Graph's definition of dataset type.
231 Returns
232 -------
233 out_dataset_type : `~lsst.daf.butler.DatasetType`
234 Dataset type seen by this task.
235 """
236 raise NotImplementedError()
238 @abstractmethod
239 def adapt_dataset_ref(self, ref: DatasetRef) -> DatasetRef:
240 """Transform the graph's definition of a dataset reference (parent
241 dataset type, with the registry or producer's storage class) to the one
242 seen by this task.
244 Parameters
245 ----------
246 ref : `~lsst.daf.butler.DatasetRef`
247 Graph's definition of the dataset reference.
249 Returns
250 -------
251 out_dataset_ref : `~lsst.daf.butler.DatasetRef`
252 Dataset reference seen by this task.
253 """
254 raise NotImplementedError()
256 def _to_xgraph_state(self) -> dict[str, Any]:
257 """Convert this edges's attributes into a dictionary suitable for use
258 in exported networkx graphs.
259 """
260 return {
261 "connection_name": self.connection_name,
262 "parent_dataset_type_name": self.parent_dataset_type_name,
263 "storage_class_name": self.storage_class_name,
264 "is_init": bool,
265 }
267 @classmethod
268 def _unreduce(cls, kwargs: dict[str, Any]) -> Self:
269 """Unpickle an `Edge` instance."""
270 return cls(**kwargs)
272 def __reduce__(self) -> tuple[Callable[[dict[str, Any]], Edge], tuple[dict[str, Any]]]:
273 return (
274 self._unreduce,
275 (
276 dict(
277 task_key=self.task_key,
278 dataset_type_key=self.dataset_type_key,
279 storage_class_name=self.storage_class_name,
280 connection_name=self.connection_name,
281 is_calibration=self.is_calibration,
282 raw_dimensions=self.raw_dimensions,
283 ),
284 ),
285 )
288class ReadEdge(Edge):
289 """Representation of an input connection (including init-inputs and
290 prerequisites) in a pipeline graph.
292 Parameters
293 ----------
294 dataset_type_key : `NodeKey`
295 Key for the dataset type node this edge is connected to. This should
296 hold the parent dataset type name for component dataset types.
297 task_key : `NodeKey`
298 Key for the task node this edge is connected to.
299 storage_class_name : `str`
300 Name of the dataset type's storage class as seen by the task.
301 connection_name : `str`
302 Internal name for the connection as seen by the task.
303 is_calibration : `bool`
304 Whether this dataset type can be included in
305 `~lsst.daf.butler.CollectionType.CALIBRATION` collections.
306 raw_dimensions : `frozenset` [ `str` ]
307 Raw dimensions from the connection definition.
308 is_prerequisite : `bool`
309 Whether this dataset must be present in the data repository prior to
310 `QuantumGraph` generation.
311 component : `str` or `None`
312 Component of the dataset type requested by the task.
313 defer_query_constraint : `bool`
314 If `True`, by default do not include this dataset type's existence as a
315 constraint on the initial data ID query in QuantumGraph generation.
317 Notes
318 -----
319 When included in an exported `networkx` graph (e.g.
320 `PipelineGraph.make_xgraph`), read edges set the following edge attributes:
322 - ``parent_dataset_type_name``
323 - ``storage_class_name``
324 - ``is_init``
325 - ``component``
326 - ``is_prerequisite``
328 As with `ReadEdge` instance attributes, these descriptions of dataset types
329 are those specific to a task, and may differ from the graph's resolved
330 dataset type or (if `PipelineGraph.resolve` has not been called) there may
331 not even be a consistent definition of the dataset type.
332 """
334 def __init__(
335 self,
336 dataset_type_key: NodeKey,
337 task_key: NodeKey,
338 *,
339 storage_class_name: str,
340 connection_name: str,
341 is_calibration: bool,
342 raw_dimensions: frozenset[str],
343 is_prerequisite: bool,
344 component: str | None,
345 defer_query_constraint: bool,
346 ):
347 super().__init__(
348 task_key=task_key,
349 dataset_type_key=dataset_type_key,
350 storage_class_name=storage_class_name,
351 connection_name=connection_name,
352 raw_dimensions=raw_dimensions,
353 is_calibration=is_calibration,
354 )
355 self.is_prerequisite = is_prerequisite
356 self.component = component
357 self.defer_query_constraint = defer_query_constraint
359 component: str | None
360 """Component to add to `parent_dataset_type_name` to form the dataset type
361 name seen by this task (`str` or `None`).
362 """
364 is_prerequisite: bool
365 """Whether this dataset must be present in the data repository prior to
366 `QuantumGraph` generation.
367 """
369 defer_query_constraint: bool
370 """If `True`, by default do not include this dataset type's existence as a
371 constraint on the initial data ID query in QuantumGraph generation.
373 This can be `True` either because the connection class had
374 ``deferQueryConstraint=True`` or because it had ``minimum=0``.
375 """
377 @property
378 def nodes(self) -> tuple[NodeKey, NodeKey]:
379 # Docstring inherited.
380 return (self.dataset_type_key, self.task_key)
382 @property
383 def dataset_type_name(self) -> str:
384 """Complete dataset type name, as seen by the task (`str`)."""
385 if self.component is not None:
386 return f"{self.parent_dataset_type_name}.{self.component}"
387 return self.parent_dataset_type_name
389 def diff(self: ReadEdge, other: ReadEdge, connection_type: str = "connection") -> list[str]:
390 # Docstring inherited.
391 result = super().diff(other, connection_type)
392 if self.defer_query_constraint != other.defer_query_constraint:
393 result.append(
394 f"{connection_type.capitalize()} {self.connection_name!r} is marked as a deferred query "
395 f"constraint {'in A but not in B' if self.defer_query_constraint else 'in B but not in A'}."
396 )
397 return result
399 def adapt_dataset_type(self, dataset_type: DatasetType) -> DatasetType:
400 # Docstring inherited.
401 if self.component is not None:
402 dataset_type = dataset_type.makeComponentDatasetType(self.component)
403 if self.storage_class_name != dataset_type.storageClass_name:
404 return dataset_type.overrideStorageClass(self.storage_class_name)
405 return dataset_type
407 def adapt_dataset_ref(self, ref: DatasetRef) -> DatasetRef:
408 # Docstring inherited.
409 if self.component is not None:
410 ref = ref.makeComponentRef(self.component)
411 if self.storage_class_name != ref.datasetType.storageClass_name:
412 return ref.overrideStorageClass(self.storage_class_name)
413 return ref
415 @classmethod
416 def _from_connection_map(
417 cls,
418 task_key: NodeKey,
419 connection_name: str,
420 connection_map: Mapping[str, BaseConnection],
421 is_prerequisite: bool = False,
422 ) -> ReadEdge:
423 """Construct a `ReadEdge` instance from a `.BaseConnection` object.
425 Parameters
426 ----------
427 task_key : `NodeKey`
428 Key for the associated task node or task init node.
429 connection_name : `str`
430 Internal name for the connection as seen by the task,.
431 connection_map : Mapping [ `str`, `.BaseConnection` ]
432 Mapping of post-configuration object to draw dataset type
433 information from, keyed by connection name.
434 is_prerequisite : `bool`, optional
435 Whether this dataset must be present in the data repository prior
436 to `QuantumGraph` generation.
438 Returns
439 -------
440 edge : `ReadEdge`
441 New edge instance.
442 """
443 connection = connection_map[connection_name]
444 parent_dataset_type_name, component = DatasetType.splitDatasetTypeName(connection.name)
445 return cls(
446 dataset_type_key=NodeKey(NodeType.DATASET_TYPE, parent_dataset_type_name),
447 task_key=task_key,
448 component=component,
449 storage_class_name=connection.storageClass,
450 # InitInput connections don't have .isCalibration.
451 is_calibration=getattr(connection, "isCalibration", False),
452 is_prerequisite=is_prerequisite,
453 connection_name=connection_name,
454 # InitInput connections don't have a .dimensions because they
455 # always have empty dimensions.
456 raw_dimensions=frozenset(getattr(connection, "dimensions", frozenset())),
457 # PrerequisiteInput and InitInput connections don't have a
458 # .deferGraphConstraint, because they never constrain the initial
459 # data ID query.
460 defer_query_constraint=(
461 getattr(connection, "deferGraphConstraint", False) or getattr(connection, "minimum", 1) == 0
462 ),
463 )
465 def _resolve_dataset_type(
466 self,
467 *,
468 current: DatasetType | None,
469 is_initial_query_constraint: bool,
470 is_prerequisite: bool | None,
471 universe: DimensionUniverse,
472 producer: str | None,
473 consumers: Sequence[str],
474 is_registered: bool,
475 visualization_only: bool,
476 ) -> tuple[DatasetType, bool, bool]:
477 """Participate in the construction of the `DatasetTypeNode` object
478 associated with this edge.
480 Parameters
481 ----------
482 current : `lsst.daf.butler.DatasetType` or `None`
483 The current graph-wide `DatasetType`, or `None`. This will always
484 be the registry's definition of the parent dataset type, if one
485 exists. If not, it will be the dataset type definition from the
486 task in the graph that writes it, if there is one. If there is no
487 such task, this will be `None`.
488 is_initial_query_constraint : `bool`
489 Whether this dataset type is currently marked as a constraint on
490 the initial data ID query in QuantumGraph generation.
491 is_prerequisite : `bool` | None`
492 Whether this dataset type is marked as a prerequisite input in all
493 edges processed so far. `None` if this is the first edge.
494 universe : `lsst.daf.butler.DimensionUniverse`
495 Object that holds all dimension definitions.
496 producer : `str` or `None`
497 The label of the task that produces this dataset type in the
498 pipeline, or `None` if it is an overall input.
499 consumers : `Sequence` [ `str` ]
500 Labels for other consuming tasks that have already participated in
501 this dataset type's resolution.
502 is_registered : `bool`
503 Whether a registration for this dataset type was found in the
504 data repository.
505 visualization_only : `bool`
506 Resolve the graph as well as possible even when dimensions and
507 storage classes cannot really be determined. This can include
508 using the ``universe.commonSkyPix`` as the assumed dimensions of
509 connections that use the "skypix" placeholder and using "<UNKNOWN>"
510 as a storage class name (which will fail if the storage class
511 itself is ever actually loaded).
513 Returns
514 -------
515 dataset_type : `DatasetType`
516 The updated graph-wide dataset type. If ``current`` was provided,
517 this must be equal to it.
518 is_initial_query_constraint : `bool`
519 If `True`, this dataset type should be included as a constraint in
520 the initial data ID query during QuantumGraph generation; this
521 requires that ``is_initial_query_constraint`` also be `True` on
522 input.
523 is_prerequisite : `bool`
524 Whether this dataset type is marked as a prerequisite input in this
525 task and all other edges processed so far.
527 Raises
528 ------
529 MissingDatasetTypeError
530 Raised if ``current is None`` and this edge cannot define one on
531 its own.
532 IncompatibleDatasetTypeError
533 Raised if ``current is not None`` and this edge's definition is not
534 compatible with it.
535 ConnectionTypeConsistencyError
536 Raised if a prerequisite input for one task appears as a different
537 kind of connection in any other task.
538 """
539 if "skypix" in self.raw_dimensions:
540 if current is None:
541 if visualization_only:
542 dimensions = universe.conform(
543 [d if d != "skypix" else universe.commonSkyPix.name for d in self.raw_dimensions]
544 )
545 else:
546 raise MissingDatasetTypeError(
547 f"DatasetType '{self.dataset_type_name}' referenced by "
548 f"{self.task_label!r} uses 'skypix' as a dimension "
549 f"placeholder, but has not been registered with the data repository. "
550 f"Note that reference catalog names are now used as the dataset "
551 f"type name instead of 'ref_cat'."
552 )
553 else:
554 rest1 = set(universe.conform(self.raw_dimensions - {"skypix"}).names)
555 rest2 = current.dimensions.names - current.dimensions.skypix
556 if rest1 != rest2:
557 raise IncompatibleDatasetTypeError(
558 f"Non-skypix dimensions for dataset type {self.dataset_type_name} declared in "
559 f"connections ({rest1}) are inconsistent with those in "
560 f"registry's version of this dataset ({rest2})."
561 )
562 dimensions = current.dimensions
563 else:
564 dimensions = universe.conform(self.raw_dimensions)
565 is_initial_query_constraint = is_initial_query_constraint and not self.defer_query_constraint
566 if is_prerequisite is None:
567 is_prerequisite = self.is_prerequisite
568 elif is_prerequisite and not self.is_prerequisite:
569 raise ConnectionTypeConsistencyError(
570 f"Dataset type {self.parent_dataset_type_name!r} is a prerequisite input to {consumers}, "
571 f"but it is not a prerequisite to {self.task_label!r}."
572 )
573 elif not is_prerequisite and self.is_prerequisite:
574 if producer is not None:
575 raise ConnectionTypeConsistencyError(
576 f"Dataset type {self.parent_dataset_type_name!r} is a prerequisite input to "
577 f"{self.task_label}, but it is produced by {producer!r}."
578 )
579 else:
580 raise ConnectionTypeConsistencyError(
581 f"Dataset type {self.parent_dataset_type_name!r} is a prerequisite input to "
582 f"{self.task_label}, but it is a regular input to {consumers!r}."
583 )
585 def report_current_origin() -> str:
586 if is_registered:
587 return "data repository"
588 elif producer is not None:
589 return f"producing task {producer!r}"
590 else:
591 return f"consuming task(s) {consumers!r}"
593 if self.component is not None:
594 if current is None:
595 if visualization_only:
596 current = DatasetType(
597 self.parent_dataset_type_name,
598 dimensions,
599 storageClass="<UNKNOWN>",
600 isCalibration=self.is_calibration,
601 )
602 else:
603 raise MissingDatasetTypeError(
604 f"Dataset type {self.parent_dataset_type_name!r} is not registered and not produced "
605 f"by this pipeline, but it is used by task {self.task_label!r}, via component "
606 f"{self.component!r}. This pipeline cannot be resolved until the parent dataset "
607 "type is registered."
608 )
609 else:
610 try:
611 all_current_components = current.storageClass.allComponents()
612 except (KeyError, ImportError):
613 if visualization_only:
614 current = DatasetType(
615 self.parent_dataset_type_name,
616 dimensions,
617 storageClass="<UNKNOWN>",
618 isCalibration=self.is_calibration,
619 )
620 return current, is_initial_query_constraint, is_prerequisite
621 raise
622 if self.component not in all_current_components:
623 raise IncompatibleDatasetTypeError(
624 f"Dataset type {self.parent_dataset_type_name!r} has storage class "
625 f"{current.storageClass_name!r} (from {report_current_origin()}), "
626 f"which does not include component {self.component!r} "
627 f"as requested by task {self.task_label!r}."
628 )
629 # Note that we can't actually make a fully-correct DatasetType
630 # for the component the task wants, because we don't have the
631 # parent storage class.
632 current_component = all_current_components[self.component]
634 if (
635 not visualization_only
636 and current_component.name != self.storage_class_name
637 and not StorageClassFactory()
638 .getStorageClass(self.storage_class_name)
639 .can_convert(current_component)
640 ):
641 raise IncompatibleDatasetTypeError(
642 f"Dataset type '{self.parent_dataset_type_name}.{self.component}' has storage class "
643 f"{all_current_components[self.component].name!r} "
644 f"(from {report_current_origin()}), which cannot be converted to "
645 f"{self.storage_class_name!r}, as requested by task {self.task_label!r}."
646 )
647 return current, is_initial_query_constraint, is_prerequisite
648 else:
649 dataset_type = DatasetType(
650 self.parent_dataset_type_name,
651 dimensions,
652 storageClass=self.storage_class_name,
653 isCalibration=self.is_calibration,
654 )
655 if current is not None:
656 if not is_registered and producer is None:
657 # Current definition comes from another consumer; we
658 # require the dataset types to be exactly equal (not just
659 # compatible), since neither connection should take
660 # precedence.
661 if dataset_type != current:
662 raise MissingDatasetTypeError(
663 f"Definitions differ for input dataset type {self.parent_dataset_type_name!r}; "
664 f"task {self.task_label!r} has {dataset_type}, but the definition "
665 f"from {report_current_origin()} is {current}. If the storage classes are "
666 "compatible but different, registering the dataset type in the data repository "
667 "in advance will avoid this error."
668 )
669 elif not visualization_only and not dataset_type.is_compatible_with(current):
670 raise IncompatibleDatasetTypeError(
671 f"Incompatible definition for input dataset type {self.parent_dataset_type_name!r}; "
672 f"task {self.task_label!r} has {dataset_type}, but the definition "
673 f"from {report_current_origin()} is {current}."
674 )
675 return current, is_initial_query_constraint, is_prerequisite
676 else:
677 return dataset_type, is_initial_query_constraint, is_prerequisite
679 def _to_xgraph_state(self) -> dict[str, Any]:
680 # Docstring inherited.
681 result = super()._to_xgraph_state()
682 result["component"] = self.component
683 result["is_prerequisite"] = self.is_prerequisite
684 return result
686 def __reduce__(self) -> tuple[Callable[[dict[str, Any]], Edge], tuple[dict[str, Any]]]:
687 return (
688 self._unreduce,
689 (
690 dict(
691 dataset_type_key=self.dataset_type_key,
692 task_key=self.task_key,
693 storage_class_name=self.storage_class_name,
694 connection_name=self.connection_name,
695 is_calibration=self.is_calibration,
696 raw_dimensions=self.raw_dimensions,
697 is_prerequisite=self.is_prerequisite,
698 component=self.component,
699 defer_query_constraint=self.defer_query_constraint,
700 ),
701 ),
702 )
705class WriteEdge(Edge):
706 """Representation of an output connection (including init-outputs) in a
707 pipeline graph.
709 Notes
710 -----
711 When included in an exported `networkx` graph (e.g.
712 `PipelineGraph.make_xgraph`), write edges set the following edge
713 attributes:
715 - ``parent_dataset_type_name``
716 - ``storage_class_name``
717 - ``is_init``
719 As with `WRiteEdge` instance attributes, these descriptions of dataset
720 types are those specific to a task, and may differ from the graph's
721 resolved dataset type or (if `PipelineGraph.resolve` has not been called)
722 there may not even be a consistent definition of the dataset type.
723 """
725 @property
726 def nodes(self) -> tuple[NodeKey, NodeKey]:
727 # Docstring inherited.
728 return (self.task_key, self.dataset_type_key)
730 def adapt_dataset_type(self, dataset_type: DatasetType) -> DatasetType:
731 # Docstring inherited.
732 if self.storage_class_name != dataset_type.storageClass_name:
733 return dataset_type.overrideStorageClass(self.storage_class_name)
734 return dataset_type
736 def adapt_dataset_ref(self, ref: DatasetRef) -> DatasetRef:
737 # Docstring inherited.
738 if self.storage_class_name != ref.datasetType.storageClass_name:
739 return ref.overrideStorageClass(self.storage_class_name)
740 return ref
742 @classmethod
743 def _from_connection_map(
744 cls,
745 task_key: NodeKey,
746 connection_name: str,
747 connection_map: Mapping[str, BaseConnection],
748 ) -> WriteEdge:
749 """Construct a `WriteEdge` instance from a `.BaseConnection` object.
751 Parameters
752 ----------
753 task_key : `NodeKey`
754 Key for the associated task node or task init node.
755 connection_name : `str`
756 Internal name for the connection as seen by the task,.
757 connection_map : Mapping [ `str`, `.BaseConnection` ]
758 Mapping of post-configuration object to draw dataset type
759 information from, keyed by connection name.
761 Returns
762 -------
763 edge : `WriteEdge`
764 New edge instance.
765 """
766 connection = connection_map[connection_name]
767 parent_dataset_type_name, component = DatasetType.splitDatasetTypeName(connection.name)
768 if component is not None:
769 raise ValueError(
770 f"Illegal output component dataset {connection.name!r} in task {task_key.name!r}."
771 )
772 return cls(
773 task_key=task_key,
774 dataset_type_key=NodeKey(NodeType.DATASET_TYPE, parent_dataset_type_name),
775 storage_class_name=connection.storageClass,
776 connection_name=connection_name,
777 # InitOutput connections don't have .isCalibration.
778 is_calibration=getattr(connection, "isCalibration", False),
779 # InitOutput connections don't have a .dimensions because they
780 # always have empty dimensions.
781 raw_dimensions=frozenset(getattr(connection, "dimensions", frozenset())),
782 )
784 def _resolve_dataset_type(self, current: DatasetType | None, universe: DimensionUniverse) -> DatasetType:
785 """Participate in the construction of the `DatasetTypeNode` object
786 associated with this edge.
788 Parameters
789 ----------
790 current : `lsst.daf.butler.DatasetType` or `None`
791 The current graph-wide `DatasetType`, or `None`. This will always
792 be the registry's definition of the parent dataset type, if one
793 exists.
794 universe : `lsst.daf.butler.DimensionUniverse`
795 Object that holds all dimension definitions.
797 Returns
798 -------
799 dataset_type : `DatasetType`
800 A dataset type compatible with this edge. If ``current`` was
801 provided, this must be equal to it.
803 Raises
804 ------
805 IncompatibleDatasetTypeError
806 Raised if ``current is not None`` and this edge's definition is not
807 compatible with it.
808 """
809 try:
810 dimensions = universe.conform(self.raw_dimensions)
811 dataset_type = DatasetType(
812 self.parent_dataset_type_name,
813 dimensions,
814 storageClass=self.storage_class_name,
815 isCalibration=self.is_calibration,
816 )
817 except Exception as err:
818 err.add_note(f"In connection {self.connection_name!r} of task {self.task_label!r}.")
819 raise
820 if current is not None:
821 if not current.is_compatible_with(dataset_type):
822 raise IncompatibleDatasetTypeError(
823 f"Incompatible definition for output dataset type {self.parent_dataset_type_name!r}: "
824 f"task {self.task_label!r} has {dataset_type}, but data repository has {current}."
825 )
826 return current
827 else:
828 return dataset_type