soorena62
83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
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(),
|
|
}
|