Coverage for python / lsst / pipe / base / pipeline_graph / _edges.py: 28%
215 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-27 01:20 -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 `~lsst.daf.butler.DatasetType`, or `None`.
484 This will always be the registry's definition of the parent dataset
485 type, if one exists. If not, it will be the dataset type
486 definition from the task in the graph that writes it, if there is
487 one. If there is no 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 : `~collections.abc.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 : `~lsst.daf.butler.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 if visualization_only and dataset_type.dimensions == current.dimensions:
663 # Make a visualization-only ambiguous storage class
664 # "name".
665 all_storage_classes = set(current.storageClass_name.split("/"))
666 all_storage_classes.update(dataset_type.storageClass_name.split("/"))
667 current = DatasetType(
668 current.name,
669 current.dimensions,
670 "/".join(sorted(all_storage_classes)),
671 )
672 else:
673 raise MissingDatasetTypeError(
674 f"Definitions differ for input dataset type "
675 f"{self.parent_dataset_type_name!r}; task {self.task_label!r} has "
676 f"{dataset_type}, but the definition from {report_current_origin()} is "
677 f"{current}. If the storage classes are compatible but different, "
678 "registering the dataset type in the data repository in advance will avoid "
679 "this error."
680 )
681 elif not visualization_only and not dataset_type.is_compatible_with(current):
682 raise IncompatibleDatasetTypeError(
683 f"Incompatible definition for input dataset type {self.parent_dataset_type_name!r}; "
684 f"task {self.task_label!r} has {dataset_type}, but the definition "
685 f"from {report_current_origin()} is {current}."
686 )
687 return current, is_initial_query_constraint, is_prerequisite
688 else:
689 return dataset_type, is_initial_query_constraint, is_prerequisite
691 def _to_xgraph_state(self) -> dict[str, Any]:
692 # Docstring inherited.
693 result = super()._to_xgraph_state()
694 result["component"] = self.component
695 result["is_prerequisite"] = self.is_prerequisite
696 return result
698 def __reduce__(self) -> tuple[Callable[[dict[str, Any]], Edge], tuple[dict[str, Any]]]:
699 return (
700 self._unreduce,
701 (
702 dict(
703 dataset_type_key=self.dataset_type_key,
704 task_key=self.task_key,
705 storage_class_name=self.storage_class_name,
706 connection_name=self.connection_name,
707 is_calibration=self.is_calibration,
708 raw_dimensions=self.raw_dimensions,
709 is_prerequisite=self.is_prerequisite,
710 component=self.component,
711 defer_query_constraint=self.defer_query_constraint,
712 ),
713 ),
714 )
717class WriteEdge(Edge):
718 """Representation of an output connection (including init-outputs) in a
719 pipeline graph.
721 Notes
722 -----
723 When included in an exported `networkx` graph (e.g.
724 `PipelineGraph.make_xgraph`), write edges set the following edge
725 attributes:
727 - ``parent_dataset_type_name``
728 - ``storage_class_name``
729 - ``is_init``
731 As with `WRiteEdge` instance attributes, these descriptions of dataset
732 types are those specific to a task, and may differ from the graph's
733 resolved dataset type or (if `PipelineGraph.resolve` has not been called)
734 there may not even be a consistent definition of the dataset type.
735 """
737 @property
738 def nodes(self) -> tuple[NodeKey, NodeKey]:
739 # Docstring inherited.
740 return (self.task_key, self.dataset_type_key)
742 def adapt_dataset_type(self, dataset_type: DatasetType) -> DatasetType:
743 # Docstring inherited.
744 if self.storage_class_name != dataset_type.storageClass_name:
745 return dataset_type.overrideStorageClass(self.storage_class_name)
746 return dataset_type
748 def adapt_dataset_ref(self, ref: DatasetRef) -> DatasetRef:
749 # Docstring inherited.
750 if self.storage_class_name != ref.datasetType.storageClass_name:
751 return ref.overrideStorageClass(self.storage_class_name)
752 return ref
754 @classmethod
755 def _from_connection_map(
756 cls,
757 task_key: NodeKey,
758 connection_name: str,
759 connection_map: Mapping[str, BaseConnection],
760 ) -> WriteEdge:
761 """Construct a `WriteEdge` instance from a `.BaseConnection` object.
763 Parameters
764 ----------
765 task_key : `NodeKey`
766 Key for the associated task node or task init node.
767 connection_name : `str`
768 Internal name for the connection as seen by the task,.
769 connection_map : Mapping [ `str`, `.BaseConnection` ]
770 Mapping of post-configuration object to draw dataset type
771 information from, keyed by connection name.
773 Returns
774 -------
775 edge : `WriteEdge`
776 New edge instance.
777 """
778 connection = connection_map[connection_name]
779 parent_dataset_type_name, component = DatasetType.splitDatasetTypeName(connection.name)
780 if component is not None:
781 raise ValueError(
782 f"Illegal output component dataset {connection.name!r} in task {task_key.name!r}."
783 )
784 return cls(
785 task_key=task_key,
786 dataset_type_key=NodeKey(NodeType.DATASET_TYPE, parent_dataset_type_name),
787 storage_class_name=connection.storageClass,
788 connection_name=connection_name,
789 # InitOutput connections don't have .isCalibration.
790 is_calibration=getattr(connection, "isCalibration", False),
791 # InitOutput connections don't have a .dimensions because they
792 # always have empty dimensions.
793 raw_dimensions=frozenset(getattr(connection, "dimensions", frozenset())),
794 )
796 def _resolve_dataset_type(self, current: DatasetType | None, universe: DimensionUniverse) -> DatasetType:
797 """Participate in the construction of the `DatasetTypeNode` object
798 associated with this edge.
800 Parameters
801 ----------
802 current : `lsst.daf.butler.DatasetType` or `None`
803 The current graph-wide `~lsst.daf.butler.DatasetType`, or `None`.
804 This will always be the registry's definition of the parent dataset
805 type, if one exists.
806 universe : `lsst.daf.butler.DimensionUniverse`
807 Object that holds all dimension definitions.
809 Returns
810 -------
811 dataset_type : `~lsst.daf.butler.DatasetType`
812 A dataset type compatible with this edge. If ``current`` was
813 provided, this must be equal to it.
815 Raises
816 ------
817 IncompatibleDatasetTypeError
818 Raised if ``current is not None`` and this edge's definition is not
819 compatible with it.
820 """
821 try:
822 dimensions = universe.conform(self.raw_dimensions)
823 dataset_type = DatasetType(
824 self.parent_dataset_type_name,
825 dimensions,
826 storageClass=self.storage_class_name,
827 isCalibration=self.is_calibration,
828 )
829 except Exception as err:
830 err.add_note(f"In connection {self.connection_name!r} of task {self.task_label!r}.")
831 raise
832 if current is not None:
833 if not current.is_compatible_with(dataset_type):
834 raise IncompatibleDatasetTypeError(
835 f"Incompatible definition for output dataset type {self.parent_dataset_type_name!r}: "
836 f"task {self.task_label!r} has {dataset_type}, but data repository has {current}."
837 )
838 return current
839 else:
840 return dataset_type