Add All Folders

Dagster/API_R&D/api/execution_context.py

@@ -0,0 +1,82 @@
+from typing import Dict, Any
+from api.assets import Asset
+from api.io_port import InPort, OutPort
+
+
+class ExecutionContext:
+    """
+    Runtime data container for a workflow execution.
+    Stores produced assets for OutPorts and allows InPorts to retrieve them.
+    """
+
+    def __init__(self, execution_id: str | None = None):
+        self.execution_id = execution_id
+
+        # Key: OutPort, Value: Asset
+        self._data: Dict[OutPort, Asset] = {}
+
+        # Arbitrary metadata (optional)
+        self._metadata: Dict[str, Any] = {}
+
+    # Store produced data
+    def put(self, output_port: OutPort, asset: Asset):
+        """
+        Register an asset produced by an OutPort.
+        Also updates output_port.produced_asset.
+        """
+        self._data[output_port] = asset
+        output_port.produced_asset = asset
+
+    # Retrieve data
+    def get(self, port):
+        """
+        Retrieve asset for:
+        - OutPort → return its asset
+        - InPort → return asset from its connected OutPort
+        """
+        if isinstance(port, OutPort):
+            return self._data.get(port)
+
+        if isinstance(port, InPort):
+            if port.connected_output is None:
+                return None
+            return self._data.get(port.connected_output)
+
+        raise TypeError("ExecutionContext.get() expects InPort or OutPort")
+
+    # Check if data exists
+    def has_data(self, port) -> bool:
+        if isinstance(port, OutPort):
+            return port in self._data
+
+        if isinstance(port, InPort):
+            if port.connected_output is None:
+                return False
+            return port.connected_output in self._data
+
+        return False
+
+    # Metadata management
+    def set_meta(self, key: str, value: Any):
+        self._metadata[key] = value
+
+    def get_meta(self, key: str, default=None):
+        return self._metadata.get(key, default)
+
+    # Debug snapshot
+    def status_view(self):
+        """
+        Lightweight snapshot of current execution state.
+        Safe for GUI display.
+        """
+        return {
+            "execution_id": self.execution_id,
+            "assets": {
+                f"{port.parent_task.name}.{port.name}": {
+                    "dtype": port.dtype.__class__.__name__,
+                    "preserved": asset.preserved,
+                }
+                for port, asset in self._data.items()
+            },
+            "metadata": self._metadata.copy(),
+        }