Merge pull request #189 from biggraph/darabos-optional-inputs

lynxkite-core/pyproject.toml

@@ -11,3 +11,6 @@ dependencies = [
 dev = [
     "pytest",
 ]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

lynxkite-core/src/lynxkite/core/executors/one_by_one.py

@@ -131,12 +131,22 @@ async def execute(ws: workspace.Workspace, catalog: ops.Catalog, cache=None):
             for task in ts:
                 try:
                     inputs = []
+                    missing = []
                     for i in op.inputs:
                         if i.position.is_vertical():
-                            assert (n, i.name) in batch_inputs, f"{i.name} is missing"
-                            inputs.append(batch_inputs[(n, i.name)])
+                            if (n, i.name) in batch_inputs:
+                                inputs.append(batch_inputs[(n, i.name)])
+                            else:
+                                opt_type = ops.get_optional_type(i.type)
+                                if opt_type is not None:
+                                    inputs.append(None)
+                                else:
+                                    missing.append(i.name)
                         else:
                             inputs.append(task)
+                    if missing:
+                        node.publish_error(f"Missing input: {', '.join(missing)}")
+                        break
                     if cache is not None:
                         key = make_cache_key((inputs, params))
                         if key not in cache:

lynxkite-core/src/lynxkite/core/executors/simple.py

@@ -42,7 +42,11 @@ async def execute(ws: workspace.Workspace, catalog: ops.Catalog):
                 if i.name in edges and edges[i.name] in outputs:
                     inputs.append(outputs[edges[i.name]])
                 else:
-                    missing.append(i.name)
+                    opt_type = ops.get_optional_type(i.type)
+                    if opt_type is not None:
+                        inputs.append(None)
+                    else:
+                        missing.append(i.name)
             if missing:
                 node.publish_error(f"Missing input: {', '.join(missing)}")
                 continue

lynxkite-core/src/lynxkite/core/ops.py

@@ -149,6 +149,16 @@ def basic_outputs(*names):
     return {name: Output(name=name, type=None) for name in names}
 
 
+def get_optional_type(type):
+    """For a type like `int | None`, returns `int`. Returns `None` otherwise."""
+    if isinstance(type, types.UnionType):
+        match type.__args__:
+            case (types.NoneType, type):
+                return type
+            case (type, types.NoneType):
+                return type
+
+
 def _param_to_type(name, value, type):
     value = value or ""
     if type is int:
@@ -160,12 +170,9 @@ def _param_to_type(name, value, type):
     if isinstance(type, enum.EnumMeta):
         assert value in type.__members__, f"{value} is not an option for {name}."
         return type[value]
-    if isinstance(type, types.UnionType):
-        match type.__args__:
-            case (types.NoneType, type):
-                return None if value == "" else _param_to_type(name, value, type)
-            case (type, types.NoneType):
-                return None if value == "" else _param_to_type(name, value, type)
+    opt_type = get_optional_type(type)
+    if opt_type:
+        return None if value == "" else _param_to_type(name, value, opt_type)
     if isinstance(type, typeof) and issubclass(type, pydantic.BaseModel):
         try:
             return type.model_validate_json(value)

lynxkite-core/src/lynxkite/core/workspace.py

@@ -125,7 +125,7 @@ class Workspace(BaseConfig):
         return self.env in ops.EXECUTORS
 
     async def execute(self):
-        await ops.EXECUTORS[self.env](self)
+        return await ops.EXECUTORS[self.env](self)
 
     def save(self, path: str):
         """Persist the workspace to a local file in JSON format."""
@@ -201,3 +201,36 @@ class Workspace(BaseConfig):
                 if "data" not in nc:
                     nc["data"] = pycrdt.Map()
                 np._crdt = nc
+
+    def add_node(self, func):
+        """For convenience in e.g. tests."""
+        random_string = os.urandom(4).hex()
+        node = WorkspaceNode(
+            id=f"{func.__op__.name} {random_string}",
+            type=func.__op__.type,
+            data=WorkspaceNodeData(
+                title=func.__op__.name,
+                params={},
+                display=None,
+                input_metadata=None,
+                error=None,
+                status=NodeStatus.planned,
+            ),
+            position=Position(x=0, y=0),
+        )
+        self.nodes.append(node)
+        return node
+
+    def add_edge(
+        self, source: WorkspaceNode, sourceHandle: str, target: WorkspaceNode, targetHandle: str
+    ):
+        """For convenience in e.g. tests."""
+        edge = WorkspaceEdge(
+            id=f"{source.id} {sourceHandle} to {target.id} {targetHandle}",
+            source=source.id,
+            target=target.id,
+            sourceHandle=sourceHandle,
+            targetHandle=targetHandle,
+        )
+        self.edges.append(edge)
+        return edge

lynxkite-core/tests/test_one_by_one.py

@@ -0,0 +1,27 @@
+from lynxkite.core import ops, workspace
+from lynxkite.core.executors import one_by_one
+
+
+async def test_optional_inputs():
+    @ops.op("test", "one")
+    def one():
+        return 1
+
+    @ops.input_position(a="bottom", b="bottom")
+    @ops.op("test", "maybe add")
+    def maybe_add(a: list[int], b: list[int] | None = None):
+        return [a + b for a, b in zip(a, b)] if b else a
+
+    assert maybe_add.__op__.inputs == [
+        ops.Input(name="a", type=list[int], position="bottom"),
+        ops.Input(name="b", type=list[int] | None, position="bottom"),
+    ]
+    one_by_one.register("test")
+    ws = workspace.Workspace(env="test", nodes=[], edges=[])
+    a = ws.add_node(one)
+    b = ws.add_node(maybe_add)
+    outputs = await ws.execute()
+    assert b.data.error == "Missing input: a"
+    ws.add_edge(a, "output", b, "a")
+    outputs = await ws.execute()
+    assert outputs[b.id].last_result == [1]

lynxkite-core/tests/test_simple.py

@@ -0,0 +1,26 @@
+from lynxkite.core import ops, workspace
+from lynxkite.core.executors import simple
+
+
+async def test_optional_inputs():
+    @ops.op("test", "one")
+    def one():
+        return 1
+
+    @ops.op("test", "maybe add")
+    def maybe_add(a: int, b: int | None = None):
+        return a + (b or 0)
+
+    assert maybe_add.__op__.inputs == [
+        ops.Input(name="a", type=int, position="left"),
+        ops.Input(name="b", type=int | None, position="left"),
+    ]
+    simple.register("test")
+    ws = workspace.Workspace(env="test", nodes=[], edges=[])
+    a = ws.add_node(one)
+    b = ws.add_node(maybe_add)
+    await ws.execute()
+    assert b.data.error == "Missing input: a"
+    ws.add_edge(a, "output", b, "a")
+    outputs = await ws.execute()
+    assert outputs[b.id, "output"] == 1

lynxkite-graph-analytics/src/lynxkite_graph_analytics/core.py

@@ -187,6 +187,7 @@ async def execute(ws: workspace.Workspace):
                 todo.remove(id)
                 progress = True
                 await _execute_node(node, ws, catalog, outputs)
+    return outputs
 
 
 async def await_if_needed(obj):
@@ -213,10 +214,15 @@ async def _execute_node(
     # Convert inputs types to match operation signature.
     try:
         inputs = []
+        missing = []
         for p in op.inputs:
             if p.name not in input_map:
-                node.publish_error(f"Missing input: {p.name}")
-                return
+                opt_type = ops.get_optional_type(p.type)
+                if opt_type is not None:
+                    inputs.append(None)
+                else:
+                    missing.append(p.name)
+                continue
             x = input_map[p.name]
             if p.type == nx.Graph and isinstance(x, Bundle):
                 x = x.to_nx()
@@ -230,6 +236,9 @@ async def _execute_node(
             traceback.print_exc()
         node.publish_error(e)
         return
+    if missing:
+        node.publish_error(f"Missing input: {', '.join(missing)}")
+        return
     # Execute op.
     try:
         result = op(*inputs, **params)

lynxkite-graph-analytics/tests/test_lynxkite_ops.py

@@ -160,5 +160,28 @@ async def test_multiple_inputs():
     assert ws.nodes[-1].data.display is False
 
 
+async def test_optional_inputs():
+    @ops.op("test", "one")
+    def one():
+        return 1
+
+    @ops.op("test", "maybe add")
+    def maybe_add(a: int, b: int | None = None):
+        return a + (b or 0)
+
+    assert maybe_add.__op__.inputs == [
+        ops.Input(name="a", type=int, position="left"),
+        ops.Input(name="b", type=int | None, position="left"),
+    ]
+    ws = workspace.Workspace(env="test", nodes=[], edges=[])
+    a = ws.add_node(one)
+    b = ws.add_node(maybe_add)
+    await execute(ws)
+    assert b.data.error == "Missing input: a"
+    ws.add_edge(a, "output", b, "a")
+    outputs = await execute(ws)
+    assert outputs[b.id, "output"] == 1
+
+
 if __name__ == "__main__":
     pytest.main()