luminarycloud 0.22.2__py3-none-any.whl → 0.23.0__py3-none-any.whl

luminarycloud/pipelines/core.py

@@ -1,12 +1,28 @@
 # Copyright 2025 Luminary Cloud, Inc. All Rights Reserved.
+from __future__ import annotations
 from abc import ABC, abstractmethod
 from dataclasses import is_dataclass, fields
-from typing import Any, Type, TypeVar, Generic
+from typing import Any, Callable, Mapping, Type, TypeVar, Generic, TYPE_CHECKING
+from .user_code_validation import validate_portable_top_level_function
 from typing_extensions import Self
+import builtins
+import inspect
 import re
+import textwrap
 import yaml
 
 from ..pipeline_util.yaml import ensure_yamlizable
+from .flowables import (
+    PipelineOutput,
+    PipelineInput,
+    FlowableType,
+    flowable_class_to_name,
+    flowable_name_to_class,
+    FlowableIOSchema,
+)
+
+if TYPE_CHECKING:
+    from .arguments import PipelineArgValueType
 
 
 class PipelineParameterRegistry:
@@ -86,40 +102,12 @@ class PipelineParameter(ABC):
         return self.__hash__() == other.__hash__()
 
 
-class PipelineInput:
-    """
-    A named input for a Stage. Explicitly connected to a PipelineOutput.
-    """
-
-    def __init__(self, upstream_output: "PipelineOutput", owner: "Stage", name: str):
-        self.upstream_output = upstream_output
-        self.owner = owner
-        self.name = name
-
-    def _to_dict(self, id_for_stage: dict) -> dict:
-        if self.upstream_output.owner not in id_for_stage:
-            raise ValueError(
-                f"Stage {self.owner} depends on a stage, {self.upstream_output.owner}, that isn't in the Pipeline. Did you forget to add it?"
-            )
-        upstream_stage_id = id_for_stage[self.upstream_output.owner]
-        upstream_output_name = self.upstream_output.name
-        return {self.name: f"{upstream_stage_id}.{upstream_output_name}"}
-
-
-class PipelineOutput(ABC):
+class StopRun(RuntimeError):
     """
-    A named output for a Stage. Can be used to spawn any number of connected PipelineInputs.
+    Raised by RunScript code to indicate that the pipeline run should stop intentionally.
     """
 
-    def __init__(self, owner: "Stage", name: str):
-        self.owner = owner
-        self.name = name
-        self.downstream_inputs: list[PipelineInput] = []
-
-    def _spawn_input(self, owner: "Stage", name: str) -> PipelineInput:
-        input = PipelineInput(self, owner, name)
-        self.downstream_inputs.append(input)
-        return input
+    pass
 
 
 class StageInputs:
@@ -187,11 +175,52 @@ class StageOutputs(ABC):
         return inputs
 
 
+class DynamicStageOutputs(StageOutputs):
+    def __init__(self, owner: "RunScript", output_types: dict[str, FlowableType]):
+        self.owner = owner
+        self._order = list(output_types.keys())
+        self.outputs: dict[str, PipelineOutput] = {}
+        for name in self._order:
+            output_type = output_types[name]
+            output_cls = flowable_name_to_class(output_type)
+            self.outputs[name] = output_cls(owner, name)
+
+    def downstream_inputs(self) -> list[PipelineInput]:
+        inputs = []
+        for output in self.outputs.values():
+            inputs.extend(output.downstream_inputs)
+        return inputs
+
+    def __getattr__(self, name: str) -> PipelineOutput:
+        return self.outputs[name]
+
+    def __getitem__(self, key: int | str) -> PipelineOutput:
+        if isinstance(key, int):
+            name = self._order[key]
+            return self.outputs[name]
+        return self.outputs[key]
+
+    def __iter__(self):
+        return iter(self._order)
+
+    def __len__(self) -> int:
+        return len(self.outputs)
+
+    def keys(self):
+        return self.outputs.keys()
+
+    def values(self):
+        return self.outputs.values()
+
+    def items(self):
+        return self.outputs.items()
+
+
 class StageRegistry:
     def __init__(self):
         self.stages = {}
 
-    def register(self, stage_class: Type["Stage"]) -> None:
+    def register(self, stage_class: Type["StandardStage"] | Type["RunScript"]) -> None:
         self.stages[stage_class.__name__] = stage_class
 
     def get(self, stage_type_name: str) -> Type["Stage"]:
@@ -203,7 +232,7 @@ class StageRegistry:
 TOutputs = TypeVar("TOutputs", bound=StageOutputs)
 
 
-class Stage(Generic[TOutputs], ABC):
+class StandardStage(Generic[TOutputs], ABC):
     def __init__(
         self,
         stage_name: str | None,
@@ -268,7 +297,7 @@ class Stage(Generic[TOutputs], ABC):
 
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
-        Stage._registry.register(cls)
+        StandardStage._registry.register(cls)
 
     @classmethod
     def _get_subclass(cls, stage_type_name: str) -> Type["Stage"]:
@@ -281,6 +310,212 @@ class Stage(Generic[TOutputs], ABC):
         return params
 
 
+class RunScript:
+    """
+    RunScript is a stage that runs a user-provided Python function.
+
+    While you can instantiate a RunScript stage directly, the usual way to construct one is to
+    decorate a function with the `@stage` decorator.
+
+    Examples
+    --------
+    >>> @pipelines.stage(
+    ...     inputs={"geometry": read_geo.outputs.geometry},
+    ...     outputs={"geometry": pipelines.PipelineOutputGeometry},
+    ... )
+    ... def ensure_single_volume(geometry: lc.Geometry):
+    ...     _, volumes = geometry.list_entities()
+    ...     if len(volumes) != 1:
+    ...         raise pipelines.StopRun("expected exactly one volume")
+    ...     return {"geometry": geometry}
+    """
+
+    def __init__(
+        self,
+        script: Callable[..., dict[str, Any] | None] | str,
+        *,
+        stage_name: str | None = None,
+        inputs: dict[str, PipelineOutput] | None = None,
+        outputs: Mapping[str, type[PipelineOutput] | str] | None = None,
+        entrypoint: str | None = None,
+        params: dict[str, Any] | None = None,
+    ):
+        inputs = inputs or {}
+        params = params or {}
+        outputs = outputs or {}
+        overlapping = set(inputs.keys()).intersection(params.keys())
+        if overlapping:
+            overlap = ", ".join(sorted(overlapping))
+            raise ValueError(f"RunScript params and inputs cannot share names: {overlap}")
+
+        inputs_and_params = set(inputs.keys()).union(params.keys())
+        validated_source = validate_portable_top_level_function(
+            script, must_have_params=inputs_and_params, can_have_params={"context"}
+        )
+        self._stage_type_name = "RunScript"
+        self._entrypoint = entrypoint or validated_source.entrypoint
+        self._name = (
+            stage_name if stage_name is not None else self._default_stage_name(self._entrypoint)
+        )
+
+        for input_name, upstream_output in inputs.items():
+            if not isinstance(upstream_output, PipelineOutput):
+                raise TypeError(
+                    f"Input '{input_name}' must be a PipelineOutput, got {type(upstream_output).__name__}"
+                )
+
+        stage_inputs_kwargs = {
+            input_name: (PipelineOutput, upstream_output)
+            for input_name, upstream_output in inputs.items()
+        }
+        self._inputs = StageInputs(self, **stage_inputs_kwargs)
+
+        input_types = {
+            input_name: flowable_class_to_name(type(upstream_output))
+            for input_name, upstream_output in inputs.items()
+        }
+        output_flowable_types = self._normalize_output_types(outputs)
+        self._io_schema = FlowableIOSchema(
+            inputs=input_types,
+            outputs=output_flowable_types,
+        )
+
+        self.outputs = DynamicStageOutputs(self, output_flowable_types)
+
+        reserved_params = {
+            "$script": validated_source.source,
+            "$output_types": {name: ft.value for name, ft in output_flowable_types.items()},
+            "$entrypoint": self._entrypoint,
+        }
+        user_params = dict(params or {})
+        invalid_param_names = ({"context"} | reserved_params.keys()).intersection(
+            user_params.keys()
+        )
+        if invalid_param_names:
+            invalid = ", ".join(sorted(invalid_param_names))
+            raise ValueError(f"RunScript params cannot use reserved names: {invalid}")
+        overlapping_input_names = set(inputs.keys()).intersection(user_params.keys())
+        if overlapping_input_names:
+            overlap = ", ".join(sorted(overlapping_input_names))
+            raise ValueError(f"RunScript params and inputs cannot share names: {overlap}")
+        if "context" in inputs.keys():
+            raise ValueError("RunScript inputs cannot include reserved name 'context'")
+
+        self._params = reserved_params | user_params
+        ensure_yamlizable(self._params_dict()[0], "RunScript parameters")
+
+    @staticmethod
+    def _default_stage_name(entrypoint: str) -> str:
+        words = entrypoint.replace("_", " ").split()
+        if not words:
+            return "RunScript"
+        return " ".join(word.capitalize() for word in words)
+
+    @staticmethod
+    def _normalize_output_types(
+        output_types: Mapping[str, type[PipelineOutput] | str | FlowableType],
+    ) -> dict[str, FlowableType]:
+        normalized: dict[str, FlowableType] = {}
+        for name, value in output_types.items():
+            if isinstance(value, FlowableType):
+                normalized[name] = value
+            elif isinstance(value, str):
+                normalized[name] = FlowableType(value)
+            elif isinstance(value, type) and issubclass(value, PipelineOutput):
+                normalized[name] = flowable_class_to_name(value)
+            else:
+                raise TypeError(
+                    f"Output '{name}' must be a PipelineOutput subclass or flowable type string, got {value}"
+                )
+        return normalized
+
+    def is_source(self) -> bool:
+        return len(self._inputs.inputs) == 0
+
+    def inputs_dict(self) -> dict[str, tuple["Stage", str]]:
+        inputs: dict[str, tuple["Stage", str]] = {}
+        for pipeline_input in self._inputs.inputs:
+            inputs[pipeline_input.name] = (
+                pipeline_input.upstream_output.owner,
+                pipeline_input.upstream_output.name,
+            )
+        return inputs
+
+    def downstream_stages(self) -> list["Stage"]:
+        return [inp.owner for inp in self.outputs.downstream_inputs()]
+
+    def _params_dict(self) -> tuple[dict, set[PipelineParameter]]:
+        d: dict[str, Any] = {}
+        pipeline_params = set()
+        for name, value in self._params.items():
+            if hasattr(value, "_to_pipeline_dict"):
+                d[name], downstream_params = value._to_pipeline_dict()
+                for param in downstream_params:
+                    if not isinstance(param, PipelineParameter):
+                        raise ValueError(
+                            f"Expected `_to_pipeline_dict()` to only return PipelineParameters, but got {type(param)}"
+                        )
+                pipeline_params.update(downstream_params)
+            else:
+                d[name] = value
+        d = {k: v for k, v in d.items() if v is not None}
+        return d, pipeline_params
+
+    def _to_dict(self, id_for_task: dict) -> tuple[dict, set[PipelineParameter]]:
+        params, pipeline_params = self._params_dict()
+        d = {
+            "name": self._name,
+            "operator": self._stage_type_name,
+            "params": params,
+            "inputs": self._inputs._to_dict(id_for_task),
+        }
+        return d, pipeline_params
+
+    @classmethod
+    def _parse_params(cls, params: dict) -> dict:
+        return params
+
+
+def stage(
+    *,
+    inputs: dict[str, PipelineOutput] | None = None,
+    outputs: dict[str, type[PipelineOutput]] | None = None,
+    stage_name: str | None = None,
+    params: dict[str, PipelineParameter | PipelineArgValueType] | None = None,
+) -> Callable[[Callable[..., dict[str, Any] | None]], RunScript]:
+    """
+    Decorator for building a RunScript stage from a Python function.
+
+    Examples
+    --------
+    >>> @pipelines.stage(
+    ...     inputs={"geometry": read_geo.outputs.geometry},
+    ...     outputs={"geometry": pipelines.PipelineOutputGeometry},
+    ... )
+    ... def ensure_single_volume(geometry: lc.Geometry):
+    ...     _, volumes = geometry.list_entities()
+    ...     if len(volumes) != 1:
+    ...         raise pipelines.StopRun("expected exactly one volume")
+    ...     return {"geometry": geometry}
+    """
+
+    def decorator(fn: Callable[..., dict[str, Any] | None]) -> RunScript:
+        return RunScript(
+            script=fn,
+            stage_name=stage_name,
+            inputs=inputs,
+            outputs=outputs,
+            params=params,
+        )
+
+    return decorator
+
+
+StandardStage._registry.register(RunScript)
+
+Stage = StandardStage | RunScript
+
+
 class Pipeline:
     def __init__(self, stages: list[Stage]):
         self.stages = stages
@@ -350,7 +585,7 @@ class Pipeline:
 
         # first, parse the pipeline parameters...
         parsed_params = {}
-        for param_name, param_metadata in d.get("params", {}).items():
+        for param_name, param_metadata in (d.get("params") or {}).items():
             parsed_params[param_name] = PipelineParameter._get_subclass(param_metadata["type"])(
                 param_name
             )
@@ -405,7 +640,7 @@ def _parse_stage(pipeline_dict: dict, stage_id: str, all_stages: dict[str, Stage
     stage_type_name = stage_dict[
         "operator"
     ]  # TODO: change key to "stage_type" when we're ready to bump the yaml schema version
-    stage_class = Stage._get_subclass(stage_type_name)
+    stage_class = StandardStage._get_subclass(stage_type_name)
 
     parsed_inputs = {}
     for input_name, input_value in stage_dict["inputs"].items():
@@ -414,13 +649,29 @@ def _parse_stage(pipeline_dict: dict, stage_id: str, all_stages: dict[str, Stage
         source_output = getattr(source_stage.outputs, source_output_name)
         parsed_inputs[input_name] = source_output
 
-    parsed_params = stage_class._parse_params(stage_dict["params"])
-
-    stage_params = {
-        "stage_name": stage_dict["name"],
-        **parsed_params,
-        **parsed_inputs,
-    }
-    stage = stage_class(**stage_params)
+    parsed_params = stage_class._parse_params(stage_dict.get("params"))
+
+    if stage_class == RunScript:
+        user_params = parsed_params.copy()
+        script = user_params.pop("$script", None)
+        output_types = user_params.pop("$output_types", None)
+        entrypoint = user_params.pop("$entrypoint", None)
+        if script is None or output_types is None:
+            raise ValueError("RunScript stages must define both `$script` and `$output_types`")
+        stage = RunScript(
+            stage_name=stage_dict["name"],
+            script=script,
+            inputs=parsed_inputs,
+            outputs=output_types,
+            entrypoint=entrypoint,
+            params=user_params,
+        )
+    else:
+        stage_params = {
+            "stage_name": stage_dict["name"],
+            **parsed_params,
+            **parsed_inputs,
+        }
+        stage = stage_class(**stage_params)
     all_stages[stage_id] = stage
     return stage

luminarycloud/pipelines/flowables.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from abc import ABC
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Type, Mapping
+
+if TYPE_CHECKING:
+    from .core import Stage
+
+
+class PipelineInput:
+    """
+    A named input for a Stage. Explicitly connected to a PipelineOutput.
+    """
+
+    def __init__(self, upstream_output: "PipelineOutput", owner: "Stage", name: str):
+        self.upstream_output = upstream_output
+        self.owner = owner
+        self.name = name
+
+    def _to_dict(self, id_for_stage: dict) -> dict:
+        if self.upstream_output.owner not in id_for_stage:
+            raise ValueError(
+                f"Stage {self.owner} depends on a stage, {self.upstream_output.owner}, that isn't in the Pipeline. Did you forget to add it?"
+            )
+        upstream_stage_id = id_for_stage[self.upstream_output.owner]
+        upstream_output_name = self.upstream_output.name
+        return {self.name: f"{upstream_stage_id}.{upstream_output_name}"}
+
+
+class PipelineOutput(ABC):
+    """
+    A named output for a Stage. Can be used to spawn any number of connected PipelineInputs.
+    """
+
+    def __init__(self, owner: "Stage", name: str):
+        self.owner = owner
+        self.name = name
+        self.downstream_inputs: list[PipelineInput] = []
+
+    def _spawn_input(self, owner: "Stage", name: str) -> PipelineInput:
+        input = PipelineInput(self, owner, name)
+        self.downstream_inputs.append(input)
+        return input
+
+
+# Concrete PipelineOutput classes, i.e. the things that can "flow" in a Pipeline
+
+
+class PipelineOutputGeometry(PipelineOutput):
+    """A representation of a Geometry in a Pipeline."""
+
+    pass
+
+
+class PipelineOutputMesh(PipelineOutput):
+    """A representation of a Mesh in a Pipeline."""
+
+    pass
+
+
+class PipelineOutputSimulation(PipelineOutput):
+    """A representation of a Simulation in a Pipeline."""
+
+    pass
+
+
+# We don't inherit from StrEnum because that was added in Python 3.11, but we still want to support
+# older versions. Inheriting from str and Enum gives us the StrEnum-like behavior we want.
+class FlowableType(str, Enum):
+    """Canonical flowable type identifiers."""
+
+    GEOMETRY = "Geometry"
+    MESH = "Mesh"
+    SIMULATION = "Simulation"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+_FLOWABLE_NAME_TO_CLASS: dict[FlowableType, Type[PipelineOutput]] = {
+    FlowableType.GEOMETRY: PipelineOutputGeometry,
+    FlowableType.MESH: PipelineOutputMesh,
+    FlowableType.SIMULATION: PipelineOutputSimulation,
+}
+
+
+def flowable_class_to_name(output_cls: Type[PipelineOutput]) -> FlowableType:
+    """
+    Convert a PipelineOutput subclass to the canonical flowable type name used in pipeline YAML.
+    """
+    for flowable_type, cls in _FLOWABLE_NAME_TO_CLASS.items():
+        if issubclass(output_cls, cls):
+            return flowable_type
+    raise ValueError(f"Unsupported PipelineOutput subclass: {output_cls.__name__}")
+
+
+def flowable_name_to_class(name: str | FlowableType) -> Type[PipelineOutput]:
+    """
+    Convert a canonical flowable type name into the corresponding PipelineOutput subclass.
+    """
+    try:
+        flowable_type = FlowableType(name)
+    except ValueError as exc:
+        supported = ", ".join(ft.value for ft in FlowableType)
+        raise ValueError(
+            f"Unknown flowable type '{name}'. Supported types are: {supported}"
+        ) from exc
+    return _FLOWABLE_NAME_TO_CLASS[flowable_type]
+
+
+def _ensure_flowable_mapping(data: Mapping[str, FlowableType | str]) -> dict[str, FlowableType]:
+    mapping: dict[str, FlowableType] = {}
+    for name, value in data.items():
+        mapping[name] = value if isinstance(value, FlowableType) else FlowableType(value)
+    return mapping
+
+
+@dataclass(slots=True)
+class FlowableIOSchema:
+    """Typed representation of RunScript input/output schema."""
+
+    inputs: dict[str, FlowableType] = field(default_factory=dict)
+    outputs: dict[str, FlowableType] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Mapping[str, FlowableType | str]]) -> "FlowableIOSchema":
+        return cls(
+            inputs=_ensure_flowable_mapping(data["inputs"]),
+            outputs=_ensure_flowable_mapping(data["outputs"]),
+        )
+
+    def to_dict(self) -> dict[str, dict[str, str]]:
+        return {
+            "inputs": {name: flowable.value for name, flowable in self.inputs.items()},
+            "outputs": {name: flowable.value for name, flowable in self.outputs.items()},
+        }