stageflow-framework 0.1.3__tar.gz → 0.1.4__tar.gz

{stageflow_framework-0.1.3 → stageflow_framework-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stageflow-framework
-Version: 0.1.3
+Version: 0.1.4
 Summary: StageFlow: Pipeline framework for stages
 Author: lethargy
 Requires-Python: >=3.10

{stageflow_framework-0.1.3 → stageflow_framework-0.1.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "stageflow-framework"
-version = "0.1.3"
+version = "0.1.4"
 description = "StageFlow: Pipeline framework for stages"
 readme = "README.md"
 authors = [

{stageflow_framework-0.1.3 → stageflow_framework-0.1.4}/stageflow/core/stage.py

@@ -1,3 +1,4 @@
+import asyncio
 import copy
 from typing import Any, TYPE_CHECKING
 import yaml
@@ -121,14 +122,29 @@ class BaseStage:
             payload=payload or {},
         ))
 
-    async def wait_input(self, type_: str, timeout: float | None = None):
+    def _get_allowed_input(self, type_: str) -> InputSpec | None:
         if self.allowed_inputs:
             allowed = {spec.type for spec in self.allowed_inputs if spec.type}
             if allowed and type_ not in allowed:
                 raise ValueError(f"Input type '{type_}' is not allowed for stage '{self.stage_name}'")
-            matching = next((spec for spec in self.allowed_inputs if spec.type == type_), None)
-        else:
-            matching = None
+            return next((spec for spec in self.allowed_inputs if spec.type == type_), None)
+        return None
+
+    def start_wait_input(self, type_: str) -> asyncio.Future:
+        self._get_allowed_input(type_)
+        return self.session.start_wait_input(type_)
+
+    async def finish_wait_input(self, type_: str, fut: asyncio.Future, timeout: float | None = None):
+        matching = self._get_allowed_input(type_)
+        result = await self.session.finish_wait_input(type_, fut, timeout=timeout)
+        if result is None:
+            return None
+        if matching and matching.payload_schema is not None:
+            validate_schema(result.get("payload", {}), matching.payload_schema, "Input payload")
+        return result
+
+    async def wait_input(self, type_: str, timeout: float | None = None):
+        matching = self._get_allowed_input(type_)
         result = await self.session.wait_input(type_, timeout=timeout)
         if result is None:
             return None

{stageflow_framework-0.1.3 → stageflow_framework-0.1.4}/stageflow_framework.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stageflow-framework
-Version: 0.1.3
+Version: 0.1.4
 Summary: StageFlow: Pipeline framework for stages
 Author: lethargy
 Requires-Python: >=3.10

{stageflow_framework-0.1.3 → stageflow_framework-0.1.4}/stageflow_framework.egg-info/SOURCES.txt

@@ -36,5 +36,6 @@ tests/test_payload_validation.py
 tests/test_pipeline_tester.py
 tests/test_session_wait_input.py
 tests/test_snapshot.py
+tests/test_stage_wait_input.py
 tests/test_std_stages.py
 tests/test_subpipeline.py

stageflow_framework-0.1.4/tests/test_stage_wait_input.py

@@ -0,0 +1,64 @@
+import unittest
+
+from stageflow.core.context import Context
+from stageflow.core.event import InputSpec
+from stageflow.core.node import TerminalNode
+from stageflow.core.pipeline import Pipeline
+from stageflow.core.session import Session
+from stageflow.core.stage import BaseStage
+
+
+def _make_pipeline():
+    end = TerminalNode(id="end", type="terminal")
+    return Pipeline(entry="end", nodes=[end])
+
+
+class StageWaitInputTests(unittest.IsolatedAsyncioTestCase):
+    async def test_start_and_finish_wait_input_success(self):
+        session = Session(id="stage1", pipeline=_make_pipeline(), context=Context())
+
+        class WaitStage(BaseStage):
+            allowed_inputs = [InputSpec(type="user", payload_schema={"id": int})]
+
+            async def run(self):
+                return None
+
+        stage = WaitStage(stage_id="s1", config={}, arguments={}, outputs={}, session=session)
+        fut = stage.start_wait_input("user")
+        await session.input("user", {"id": 10})
+
+        result = await stage.finish_wait_input("user", fut, timeout=1)
+        self.assertEqual(result["payload"]["id"], 10)
+
+    async def test_finish_wait_input_validates_payload(self):
+        session = Session(id="stage2", pipeline=_make_pipeline(), context=Context())
+
+        class WaitStage(BaseStage):
+            allowed_inputs = [InputSpec(type="user", payload_schema={"id": int})]
+
+            async def run(self):
+                return None
+
+        stage = WaitStage(stage_id="s1", config={}, arguments={}, outputs={}, session=session)
+        fut = stage.start_wait_input("user")
+        await session.input("user", {"id": "bad"})
+
+        with self.assertRaises(ValueError):
+            await stage.finish_wait_input("user", fut, timeout=1)
+
+    async def test_start_wait_input_rejects_disallowed_type(self):
+        session = Session(id="stage3", pipeline=_make_pipeline(), context=Context())
+
+        class LimitedStage(BaseStage):
+            allowed_inputs = [InputSpec(type="only")]
+
+            async def run(self):
+                return None
+
+        stage = LimitedStage(stage_id="s1", config={}, arguments={}, outputs={}, session=session)
+        with self.assertRaises(ValueError):
+            stage.start_wait_input("other")
+
+
+if __name__ == "__main__":
+    unittest.main()