feat: add easier access to source, function and operator pods in pipeline

Author: eywalker

src/orcapod/pipeline/graph.py

@@ -4,7 +4,7 @@
 from orcapod import contexts
 from orcapod.protocols import core_protocols as cp
 from orcapod.protocols import database_protocols as dbp
-from typing import Any
+from typing import Any, cast
 from collections.abc import Collection
 import os
 import tempfile
@@ -94,11 +94,39 @@ def __init__(
             self.results_store_path_prefix = self.name + ("_results",)
         self.pipeline_database = pipeline_database
         self.results_database = results_database
-        self.nodes: dict[str, Node] = {}
+        self._nodes: dict[str, Node] = {}
         self.auto_compile = auto_compile
         self._dirty = False
         self._ordered_nodes = []  # Track order of invocations
 
+    @property
+    def nodes(self) -> dict[str, Node]:
+        return self._nodes.copy()
+
+    @property
+    def function_pods(self) -> dict[str, cp.Pod]:
+        return {
+            label: cast(cp.Pod, node)
+            for label, node in self._nodes.items()
+            if getattr(node, "kernel_type") == "function"
+        }
+
+    @property
+    def source_pods(self) -> dict[str, cp.Source]:
+        return {
+            label: node
+            for label, node in self._nodes.items()
+            if getattr(node, "kernel_type") == "source"
+        }
+
+    @property
+    def operator_pods(self) -> dict[str, cp.Kernel]:
+        return {
+            label: node
+            for label, node in self._nodes.items()
+            if getattr(node, "kernel_type") == "operator"
+        }
+
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         """
         Exit the pipeline context, ensuring all nodes are properly closed.
@@ -156,13 +184,13 @@ def compile(self) -> None:
                 # If there are multiple nodes with the same label, we need to resolve the collision
                 logger.info(f"Collision detected for label '{label}': {nodes}")
                 for i, node in enumerate(nodes, start=1):
-                    self.nodes[f"{label}_{i}"] = node
+                    self._nodes[f"{label}_{i}"] = node
                     node.label = f"{label}_{i}"
             else:
-                self.nodes[label] = nodes[0]
+                self._nodes[label] = nodes[0]
                 nodes[0].label = label
 
-        self.label_lut = {v: k for k, v in self.nodes.items()}
+        self.label_lut = {v: k for k, v in self._nodes.items()}
 
         self.graph = node_graph
 
@@ -172,7 +200,7 @@ def show_graph(self, **kwargs) -> None:
     def set_mode(self, mode: str) -> None:
         if mode not in ("production", "development"):
             raise ValueError("Mode must be either 'production' or 'development'")
-        for node in self.nodes.values():
+        for node in self._nodes.values():
             if hasattr(node, "set_mode"):
                 node.set_mode(mode)
 
@@ -257,27 +285,27 @@ def wrap_invocation(
 
     def __getattr__(self, item: str) -> Any:
         """Allow direct access to pipeline attributes."""
-        if item in self.nodes:
-            return self.nodes[item]
+        if item in self._nodes:
+            return self._nodes[item]
         raise AttributeError(f"Pipeline has no attribute '{item}'")
 
     def __dir__(self) -> list[str]:
         """Return a list of attributes and methods of the pipeline."""
-        return list(super().__dir__()) + list(self.nodes.keys())
+        return list(super().__dir__()) + list(self._nodes.keys())
 
     def rename(self, old_name: str, new_name: str) -> None:
         """
         Rename a node in the pipeline.
         This will update the label and the internal mapping.
         """
-        if old_name not in self.nodes:
+        if old_name not in self._nodes:
             raise KeyError(f"Node '{old_name}' does not exist in the pipeline.")
-        if new_name in self.nodes:
+        if new_name in self._nodes:
             raise KeyError(f"Node '{new_name}' already exists in the pipeline.")
-        node = self.nodes[old_name]
-        del self.nodes[old_name]
+        node = self._nodes[old_name]
+        del self._nodes[old_name]
         node.label = new_name
-        self.nodes[new_name] = node
+        self._nodes[new_name] = node
         logger.info(f"Node '{old_name}' renamed to '{new_name}'")