lsst-pipe-base 29.2025.1400__py3-none-any.whl → 29.2025.1600__py3-none-any.whl

lsst/pipe/base/pipeline_graph/visualization/_mermaid.py

@@ -32,8 +32,8 @@ import html
 import os
 import sys
 from collections.abc import Mapping
-from io import BufferedIOBase, BytesIO, StringIO, TextIOBase
-from typing import Any, TextIO
+from io import StringIO
+from typing import IO, Any
 
 from .._nodes import NodeType
 from .._pipeline_graph import PipelineGraph
@@ -58,7 +58,7 @@ _OVERFLOW_MAX_LINES = 20
 
 def show_mermaid(
     pipeline_graph: PipelineGraph,
-    stream: TextIO | BytesIO = sys.stdout,
+    stream: IO[Any] = sys.stdout,
     output_format: str = "mmd",
     width: int | None = None,
     height: int | None = None,
@@ -78,7 +78,7 @@ def show_mermaid(
     ----------
     pipeline_graph : `PipelineGraph`
         The pipeline graph to visualize.
-    stream : `TextIO` or `BytesIO`, optional
+    stream : `typing.IO`, optional
         The output stream where Mermaid code is written. Defaults to
         `sys.stdout`.
     output_format : str, optional
@@ -113,19 +113,11 @@ def show_mermaid(
     mermaid_source = _generate_mermaid_source(pipeline_graph, **kwargs)
 
     if output_format == "mmd":
-        if isinstance(stream, TextIOBase):
-            # Write Mermaid source as a string.
-            stream.write(mermaid_source)
-        else:
-            raise TypeError(f"Expected a text stream, but got {type(stream)}.")
+        # Write Mermaid source as a string.
+        stream.write(mermaid_source)
     else:
-        if isinstance(stream, BufferedIOBase):
-            # Render Mermaid source as an image and write to binary stream.
-            _render_mermaid_image(
-                mermaid_source, stream, output_format, width=width, height=height, scale=scale
-            )
-        else:
-            raise ValueError(f"Expected a binary stream, but got {type(stream)}.")
+        # Render Mermaid source as an image and write to binary stream.
+        _render_mermaid_image(mermaid_source, stream, output_format, width=width, height=height, scale=scale)
 
 
 def _generate_mermaid_source(pipeline_graph: PipelineGraph, **kwargs: Any) -> str:
@@ -210,7 +202,7 @@ def _generate_mermaid_source(pipeline_graph: PipelineGraph, **kwargs: Any) -> st
 
 def _render_mermaid_image(
     mermaid_source: str,
-    binary_stream: BytesIO,
+    binary_stream: IO[bytes],
     output_format: str,
     width: int | None = None,
     height: int | None = None,
@@ -287,7 +279,7 @@ def _render_task_node(
     node_key: NodeKey,
     node_data: Mapping[str, Any],
     options: NodeAttributeOptions,
-    stream: TextIO,
+    stream: IO[str],
 ) -> None:
     """Render a Mermaid node for a task or task-init node.
 
@@ -301,7 +293,7 @@ def _render_task_node(
     options : NodeAttributeOptions
         Rendering options controlling whether to show dimensions, storage
         classes, etc.
-    stream : TextIO
+    stream : `typing.IO` [ `str` ]
         The output stream for Mermaid syntax.
     """
     # Convert node_key into a label, handling line splitting and prefix
@@ -337,7 +329,7 @@ def _render_dataset_type_node(
     node_key: NodeKey,
     node_data: Mapping[str, Any],
     options: NodeAttributeOptions,
-    stream: TextIO,
+    stream: IO[str],
     overflow_ref: int,
 ) -> tuple[int, list[str]]:
     """Render a Mermaid node for a dataset-type node, handling overflow lines
@@ -355,7 +347,7 @@ def _render_dataset_type_node(
     options : NodeAttributeOptions
         Rendering options controlling whether to show dimensions and storage
         classes.
-    stream : TextIO
+    stream : `typing.IO` [ `str` ]
         The output stream for Mermaid syntax.
     overflow_ref : int
         The current reference number for overflow nodes. If overflow occurs,
@@ -414,7 +406,7 @@ def _render_dataset_type_node(
     return overflow_ref, overflow_ids
 
 
-def _render_simple_node(node_id: str, lines: list[str], node_class: str, stream: TextIO) -> None:
+def _render_simple_node(node_id: str, lines: list[str], node_class: str, stream: IO[str]) -> None:
     """Render a simple Mermaid node with given lines and a class.
 
     This helper function is used for both primary nodes and overflow nodes once
@@ -429,7 +421,7 @@ def _render_simple_node(node_id: str, lines: list[str], node_class: str, stream:
     node_class : str
         Mermaid class name to style the node (e.g., 'dsType', 'task',
         'taskInit').
-    stream : TextIO
+    stream : `typing.IO` [ `str` ]
         The output stream.
     """
     label = "<br>".join(lines)
@@ -437,7 +429,7 @@ def _render_simple_node(node_id: str, lines: list[str], node_class: str, stream:
     print(f"class {node_id} {node_class};", file=stream)
 
 
-def _render_edge(from_node_id: str, to_node_id: str, is_prerequisite: bool, stream: TextIO) -> None:
+def _render_edge(from_node_id: str, to_node_id: str, is_prerequisite: bool, stream: IO[str]) -> None:
     """Render a Mermaid edge from one node to another.
 
     Edges in Mermaid are normally specified as `A --> B`. Prerequisite edges
@@ -453,7 +445,7 @@ def _render_edge(from_node_id: str, to_node_id: str, is_prerequisite: bool, stre
     is_prerequisite : bool
         If True, this edge represents a prerequisite connection and will be
         styled as dashed.
-    stream : TextIO
+    stream : `typing.IO` [ `str` ]
         The output stream for Mermaid syntax.
     """
     # At this stage, we simply print the edge. The styling (dashed) for

lsst/pipe/base/pipeline_graph/visualization/_options.py

@@ -32,6 +32,7 @@ import dataclasses
 from typing import Literal
 
 from .._nodes import NodeType
+from ._status_annotator import NodeStatusOptions
 
 
 @dataclasses.dataclass
@@ -71,8 +72,8 @@ class NodeAttributeOptions:
     - `None`: context-dependent default behavior.
     """
 
-    def __bool__(self) -> bool:
-        return bool(self.dimensions or self.storage_classes or self.task_classes)
+    status: NodeStatusOptions | None
+    """Options for displaying execution status."""
 
     def has_details(self, node_type: NodeType) -> bool:
         """Check whether there is any information beyond the node name for a
@@ -93,7 +94,10 @@ class NodeAttributeOptions:
         else:
             return bool(self.dimensions or self.task_classes)
 
-    def checked(self, is_resolved: bool) -> NodeAttributeOptions:
+    def __bool__(self) -> bool:
+        return bool(self.dimensions or self.storage_classes or self.task_classes or self.status)
+
+    def checked(self, is_resolved: bool, has_status: bool = False) -> NodeAttributeOptions:
         """Check these options against a pipeline graph's resolution status and
         fill in defaults.
 
@@ -102,6 +106,9 @@ class NodeAttributeOptions:
         is_resolved : `bool`
             Whether the pipeline graph to be displayed is resolved
             (`PipelineGraph.is_fully_resolved`).
+        has_status : `bool`
+            Whether the pipeline graph to be displayed has status information.
+            Defaults to `False`.
 
         Returns
         -------
@@ -127,4 +134,5 @@ class NodeAttributeOptions:
                 self.task_classes if self.task_classes is not None else ("concise" if is_resolved else False)
             ),
             storage_classes=(self.storage_classes if self.storage_classes is not None else is_resolved),
+            status=self.status if has_status else None,
         )

lsst/pipe/base/pipeline_graph/visualization/_show.py

@@ -48,6 +48,7 @@ from ._merge import (
 )
 from ._options import NodeAttributeOptions
 from ._printer import make_default_printer
+from ._status_annotator import NodeStatusAnnotator, NodeStatusOptions
 
 DisplayNodeKey = NodeKey | MergedNodeKey
 
@@ -64,6 +65,8 @@ def parse_display_args(
     merge_output_trees: int = 4,
     merge_intermediates: bool = True,
     include_automatic_connections: bool = False,
+    status_annotator: NodeStatusAnnotator | None = None,
+    status_options: NodeStatusOptions | None = None,
 ) -> tuple[networkx.DiGraph | networkx.MultiDiGraph, NodeAttributeOptions]:
     """Print a text-based ~.PipelineGraph` visualization.
 
@@ -126,21 +129,34 @@ def parse_display_args(
     include_automatic_connections : `bool`, optional
         Whether to include automatically-added connections like the config,
         log, and metadata dataset types for each task.  Default is `False`.
+    status_annotator : `NodeStatusAnnotator`, optional
+        Annotator to add status information to the graph.  Default is `None`.
+    status_options : `NodeStatusOptions`, optional
+        Options for displaying execution status.  Default is `None`.
     """
     if init is None:
         if not dataset_types:
             raise ValueError("Cannot show init and runtime graphs unless dataset types are shown.")
         xgraph = pipeline_graph.make_xgraph()
+        if status_annotator is not None:
+            raise ValueError("Cannot show status with both init and runtime graphs.")
     elif dataset_types:
         xgraph = pipeline_graph.make_bipartite_xgraph(init)
+        if status_annotator is not None:
+            status_annotator(xgraph, dataset_types=True)
     else:
         xgraph = pipeline_graph.make_task_xgraph(init)
         storage_classes = False
+        if status_annotator is not None:
+            status_annotator(xgraph, dataset_types=False)
 
     options = NodeAttributeOptions(
-        dimensions=dimensions, storage_classes=storage_classes, task_classes=task_classes
+        dimensions=dimensions,
+        storage_classes=storage_classes,
+        task_classes=task_classes,
+        status=status_options,
     )
-    options = options.checked(pipeline_graph.is_fully_resolved)
+    options = options.checked(pipeline_graph.is_fully_resolved, has_status=status_annotator is not None)
 
     if dataset_types and not include_automatic_connections:
         taskish_nodes: list[TaskNode | TaskInitNode] = []
@@ -219,10 +235,14 @@ def show(
     if width < 0:
         width, _ = get_terminal_size()
 
+    # Number of columns used for padding after a symbol.
+    # Must match the padding added by the `Printer` class.
+    symbol_padding = 2
+
     printer = make_default_printer(layout.width, color, stream)
     printer.get_symbol = get_node_symbol
 
-    get_text = GetNodeText(xgraph, options, (width - printer.width) if width else 0)
+    get_text = GetNodeText(xgraph, options, (width - printer.width - symbol_padding) if width else 0)
     printer.get_text = get_text
 
     printer.print(stream, layout)

lsst/pipe/base/pipeline_graph/visualization/_status_annotator.py

@@ -0,0 +1,250 @@
+# This file is part of pipe_base.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (http://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This software is dual licensed under the GNU General Public License and also
+# under a 3-clause BSD license. Recipients may choose which of these licenses
+# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
+# respectively.  If you choose the GPL option then the following text applies
+# (but note that there is still no warranty even if you opt for BSD instead):
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+from __future__ import annotations
+
+__all__ = (
+    "QuantumGraphExecutionStatusAnnotator",
+    "QuantumGraphExecutionStatusOptions",
+    "QuantumProvenanceGraphStatusAnnotator",
+    "QuantumProvenanceGraphStatusOptions",
+)
+
+import dataclasses
+from typing import TYPE_CHECKING, Any, Literal, Protocol, overload
+
+import networkx
+
+from .._nodes import NodeKey, NodeType
+
+if TYPE_CHECKING:
+    from ... import quantum_provenance_graph as qpg
+
+# ANSI color codes.
+GREEN = "\033[32m"
+RED = "\033[31m"
+YELLOW = "\033[33m"
+CYAN = "\033[36m"
+WHITE = "\033[37m"
+GRAY = "\033[90m"
+MAGENTA = "\033[35m"
+BROWN = "\u001b[38;5;130m"
+RESET = "\033[0m"
+
+
+@dataclasses.dataclass
+class TaskStatusInfo:
+    """Holds status information for a task."""
+
+    expected: int
+    succeeded: int
+    failed: int
+    blocked: int
+    ready: int | None = None
+    running: int | None = None
+    wonky: int | None = None
+    unknown: int | None = None
+
+
+@dataclasses.dataclass
+class DatasetTypeStatusInfo:
+    """Holds status information for a dataset type."""
+
+    expected: int
+    produced: int
+
+
+@dataclasses.dataclass
+class StatusColors:
+    """Base class for holding ANSI color codes for different progress segments
+    or statuses.
+    """
+
+    # Base task status colors.
+    expected: str = WHITE
+    succeeded: str = GREEN
+    failed: str = RED
+
+    # Base dataset type status colors.
+    produced: str = GREEN
+
+    # Reset to default color.
+    reset: str = RESET
+
+
+@dataclasses.dataclass
+class QuantumGraphExecutionStatusColors(StatusColors):
+    """Holds ANSI color codes for different progress segments or statuses for
+    quantum graph execution reports.
+
+    Status colors for both task and dataset type nodes are included.
+    """
+
+    def __post_init__(self) -> None:
+        raise NotImplementedError("`QuantumGraphExecutionStatusColors` is not implemented yet.")
+
+
+@dataclasses.dataclass
+class QuantumProvenanceGraphStatusColors(StatusColors):
+    """Holds ANSI color codes for different progress segments or statuses for
+    quantum provenance graph reports.
+
+    Status colors for both task and dataset type nodes are included.
+    """
+
+    # Additional task status colors.
+    blocked: str = YELLOW
+    ready: str = GRAY
+    running: str = MAGENTA
+    wonky: str = CYAN
+    unknown: str = BROWN
+
+
+@dataclasses.dataclass
+class NodeStatusOptions:
+    """Base options for node status visualization.
+
+    Attributes
+    ----------
+    colors : `StatusColors`
+        A dataclass specifying ANSI color codes for distinct progress segments
+        or statuses.
+    display_percent : `bool`
+        Whether to show percentage of progress.
+    display_counts : `bool`
+        Whether to show numeric counts (e.g., succeeded/expected).
+    visualize : `bool`
+        If `True`, status information for task or dataset type nodes will be
+        visually indicated by segmented fills in text-based bars or flowchart
+        nodes.
+    min_bar_width : `int`
+        Minimum width of the visualized progress bar in characters. Only counts
+        the width of the bar itself, not any surrounding text. Only relevant if
+        `visualize` is `True` and it's a text-based visualization.
+    abbreviate : `bool`
+        If `True`, status labels will be abbreviated to save space. For
+        example, 'expected' will be abbreviated to 'exp' and 'blocked' to
+        'blk'.
+    """
+
+    colors: QuantumGraphExecutionStatusColors | QuantumProvenanceGraphStatusColors
+    display_percent: bool = True
+    display_counts: bool = True
+    visualize: bool = True
+    min_bar_width: int = 15
+    abbreviate: bool = True
+
+    def __post_init__(self) -> None:
+        if not (self.display_percent or self.display_counts or self.visualize):
+            raise ValueError(
+                "At least one of 'display_percent', 'display_counts', or 'visualize' must be True."
+            )
+
+
+@dataclasses.dataclass
+class QuantumGraphExecutionStatusOptions(NodeStatusOptions):
+    """Specialized status options for quantum graph execution reports."""
+
+    colors: QuantumGraphExecutionStatusColors = dataclasses.field(
+        default_factory=QuantumGraphExecutionStatusColors
+    )
+
+
+@dataclasses.dataclass
+class QuantumProvenanceGraphStatusOptions(NodeStatusOptions):
+    """Specialized status options for quantum provenance graph reports."""
+
+    colors: QuantumProvenanceGraphStatusColors = dataclasses.field(
+        default_factory=QuantumProvenanceGraphStatusColors
+    )
+
+
+class NodeStatusAnnotator(Protocol):
+    """Protocol for annotating a networkx graph with task and dataset type
+    status information.
+    """
+
+    @overload
+    def __call__(self, xgraph: networkx.DiGraph, dataset_types: Literal[False]) -> None: ...
+
+    @overload
+    def __call__(self, xgraph: networkx.MultiDiGraph, dataset_types: Literal[True]) -> None: ...
+
+    def __call__(self, xgraph: networkx.DiGraph | networkx.MultiDiGraph, dataset_types: bool) -> None: ...
+
+
+class QuantumGraphExecutionStatusAnnotator:
+    """Annotates a networkx graph with task and dataset status information from
+    a quantum graph execution summary, implementing the StatusAnnotator
+    protocol to update the graph with status data.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        raise NotImplementedError("`QuantumGraphExecutionStatusAnnotator` is not implemented yet.")
+
+
+class QuantumProvenanceGraphStatusAnnotator:
+    """Annotates a networkx graph with task and dataset status information from
+    a quantum provenance summary, implementing the StatusAnnotator protocol to
+    update the graph with status data.
+
+    Parameters
+    ----------
+    qpg_summary : `~lsst.pipe.base.quantum_provenance_graph.Summary`
+        The quantum provenance summary to use for status information.
+    """
+
+    def __init__(self, qpg_summary: qpg.Summary) -> None:
+        self.qpg_summary = qpg_summary
+
+    @overload
+    def __call__(self, xgraph: networkx.DiGraph, dataset_types: Literal[False]) -> None: ...
+
+    @overload
+    def __call__(self, xgraph: networkx.MultiDiGraph, dataset_types: Literal[True]) -> None: ...
+
+    def __call__(self, xgraph: networkx.DiGraph | networkx.MultiDiGraph, dataset_types: bool) -> None:
+        for task_label, task_summary in self.qpg_summary.tasks.items():
+            fields = {
+                name.replace("n_", "").replace("successful", "succeeded"): getattr(task_summary, name)
+                for name in dir(task_summary)
+                if name.startswith("n_")
+            }
+            assert sum(fields.values()) == 2 * task_summary.n_expected, f"Incosistent status counts: {fields}"
+            task_status_info = TaskStatusInfo(**fields)
+
+            key = NodeKey(NodeType.TASK, task_label)
+            xgraph.nodes[key]["status"] = task_status_info
+
+        if dataset_types:
+            for dataset_type_name, dataset_type_summary in self.qpg_summary.datasets.items():
+                expected = dataset_type_summary.n_expected
+                produced = dataset_type_summary.n_visible + dataset_type_summary.n_shadowed
+                assert produced <= expected, f"Dataset types produced ({produced}) > expected ({expected})"
+                dataset_type_status_info = DatasetTypeStatusInfo(expected=expected, produced=produced)
+
+                key = NodeKey(NodeType.DATASET_TYPE, dataset_type_name)
+                xgraph.nodes[key]["status"] = dataset_type_status_info

lsst/pipe/base/quantum_provenance_graph.py

@@ -2106,6 +2106,12 @@ class _QuantumBackedButlerFactory:
 def _cli() -> None:
     import argparse
 
+    from .pipeline_graph.visualization import (
+        QuantumProvenanceGraphStatusAnnotator,
+        QuantumProvenanceGraphStatusOptions,
+        show,
+    )
+
     parser = argparse.ArgumentParser(
         "QuantumProvenanceGraph command-line utilities.",
         description=(
@@ -2114,6 +2120,7 @@ def _cli() -> None:
         ),
     )
     subparsers = parser.add_subparsers(dest="cmd")
+
     pprint_parser = subparsers.add_parser("pprint", help="Print a saved summary as a series of tables.")
     pprint_parser.add_argument("file", type=argparse.FileType("r"), help="Saved summary JSON file.")
     pprint_parser.add_argument(
@@ -2127,12 +2134,33 @@ def _cli() -> None:
         action=argparse.BooleanOptionalAction,
         default=True,
     )
+
+    xgraph_parser = subparsers.add_parser("xgraph", help="Print a visual representation of a saved xgraph.")
+    xgraph_parser.add_argument("file", type=argparse.FileType("r"), help="Saved summary JSON file.")
+    xgraph_parser.add_argument("qgraph", type=str, help="Saved quantum graph file.")
+
     args = parser.parse_args()
+
     match args.cmd:
         case "pprint":
             summary = Summary.model_validate_json(args.file.read())
             args.file.close()
             summary.pprint(brief=args.brief, datasets=args.datasets)
+        case "xgraph":
+            summary = Summary.model_validate_json(args.file.read())
+            args.file.close()
+            status_annotator = QuantumProvenanceGraphStatusAnnotator(summary)
+            status_options = QuantumProvenanceGraphStatusOptions(
+                display_percent=True, display_counts=True, abbreviate=True, visualize=True
+            )
+            qgraph = QuantumGraph.loadUri(args.qgraph)
+            pipeline_graph = qgraph.pipeline_graph
+            show(
+                pipeline_graph,
+                dataset_types=True,
+                status_annotator=status_annotator,
+                status_options=status_options,
+            )
         case _:
             raise AssertionError(f"Unhandled subcommand {args.dest}.")
 

lsst/pipe/base/version.py

@@ -1,2 +1,2 @@
 __all__ = ["__version__"]
-__version__ = "29.2025.1400"
+__version__ = "29.2025.1600"

{lsst_pipe_base-29.2025.1400.dist-info → lsst_pipe_base-29.2025.1600.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lsst-pipe-base
-Version: 29.2025.1400
+Version: 29.2025.1600
 Summary: Pipeline infrastructure for the Rubin Science Pipelines.
 Author-email: Rubin Observatory Data Management <dm-admin@lists.lsst.org>
 License: BSD 3-Clause License
@@ -27,6 +27,7 @@ Requires-Dist: lsst-pex-config
 Requires-Dist: astropy
 Requires-Dist: pydantic<3.0,>=2
 Requires-Dist: networkx
+Requires-Dist: wcwidth
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: numpy>=1.17
 Requires-Dist: frozendict