ordeq: extend type hint for node factories

Author: nielsneerhoff

packages/ordeq/src/ordeq/_nodes.py

@@ -298,7 +298,7 @@ def create_node(
     func: Callable[FuncParams, FuncReturns],
     *,
     inputs: Sequence[Input | Node] | Input | Node | None = None,
-    outputs: Sequence[Output] | Output | None = None,
+    outputs: Sequence[Output] | Output,
     checks: Sequence[Input | Output | Node]
     | Input
     | Output
@@ -435,15 +435,15 @@ def create_node(
 def _not_passed(*args, **kwargs): ...
 
 
-not_passed = cast("Node", _not_passed)
+not_passed = cast("View", _not_passed)
 
 
 @overload
 def node(
     func: Callable[FuncParams, FuncReturns],
     *,
-    inputs: Sequence[Input | Node] | Input | Node | None = None,
-    outputs: Sequence[Output] | Output | None = None,
+    inputs: Sequence[Input | View] | Input | View | None = None,
+    outputs: Sequence[Output] | Output,
     checks: Sequence[Input | Output | Node]
     | Input
     | Output
@@ -456,9 +456,25 @@ def node(
 
 @overload
 def node(
+    func: Callable[FuncParams, FuncReturns],
     *,
-    inputs: Sequence[Input | Node] | Input | Node = not_passed,
-    outputs: Sequence[Output] | Output | None = None,
+    inputs: Sequence[Input | View] | Input | View | None = None,
+    outputs: None = None,
+    checks: Sequence[Input | Output | Node]
+    | Input
+    | Output
+    | Node
+    | ResourceType
+    | None = None,
+    **attributes: Any,
+) -> View[FuncParams, FuncReturns]: ...
+
+
+@overload
+def node(
+    *,
+    inputs: Sequence[Input | View] | Input | View = not_passed,
+    outputs: Sequence[Output] | Output,
     checks: Sequence[Input | Output | Node]
     | Input
     | Output
@@ -471,10 +487,27 @@ def node(
 ]: ...
 
 
+@overload
 def node(
-    func: Callable[FuncParams, FuncReturns] = not_passed,
     *,
-    inputs: Sequence[Input | Node] | Input | Node | None = None,
+    inputs: Sequence[Input | View] | Input | View = not_passed,
+    outputs: None = None,
+    checks: Sequence[Input | Output | Node]
+    | Input
+    | Output
+    | Node
+    | ResourceType
+    | None = None,
+    **attributes: Any,
+) -> Callable[
+    [Callable[FuncParams, FuncReturns]], View[FuncParams, FuncReturns]
+]: ...
+
+
+def node(
+    func: Callable[FuncParams, FuncReturns] = _not_passed,
+    *,
+    inputs: Sequence[Input | View] | Input | View | None = None,
     outputs: Sequence[Output] | Output | None = None,
     checks: Sequence[Input | Output | Node]
     | Input
@@ -485,9 +518,11 @@ def node(
     **attributes: Any,
 ) -> (
     Callable[
-        [Callable[FuncParams, FuncReturns]], Node[FuncParams, FuncReturns]
+        [Callable[FuncParams, FuncReturns]],
+        Node[FuncParams, FuncReturns] | View[FuncParams, FuncReturns],
     ]
     | Node[FuncParams, FuncReturns]
+    | View[FuncParams, FuncReturns]
 ):
     """Decorator that creates a node from a function. When a node is run,
     the inputs are loaded and passed to the function. The returned values

packages/ordeq/tests/resources/nodes/node_type_hints.py

@@ -1,4 +1,5 @@
 import inspect
+from typing_extensions import reveal_type
 
 from ordeq import node
 from ordeq._nodes import _is_node
@@ -13,6 +14,7 @@ def func(x: str, y: str) -> tuple[str, str]:
     return f"{x} + {y}", y
 
 
+reveal_type(func)
 print(type(func))
 print(func)
 print(inspect.get_annotations(func))

packages/ordeq/tests/resources/views/view_type_hints.py

@@ -0,0 +1,18 @@
+import inspect
+from typing_extensions import reveal_type
+
+from ordeq import node
+from ordeq._nodes import _is_node
+from ordeq_common import StringBuffer
+
+
+@node(inputs=[StringBuffer("x"), StringBuffer("y")])
+def func(x: str, y: str) -> tuple[str, str]:
+    return f"{x} + {y}", y
+
+
+reveal_type(func)
+print(type(func))
+print(func)
+print(inspect.get_annotations(func))
+print(_is_node(func))

packages/ordeq/tests/snapshots/nodes/node_type_hints.snapshot.md

@@ -2,6 +2,7 @@
 
 ```python
 import inspect
+from typing_extensions import reveal_type
 
 from ordeq import node
 from ordeq._nodes import _is_node
@@ -16,6 +17,7 @@ def func(x: str, y: str) -> tuple[str, str]:
     return f"{x} + {y}", y
 
 
+reveal_type(func)
 print(type(func))
 print(func)
 print(inspect.get_annotations(func))
@@ -31,4 +33,11 @@ node 'func' in module '__main__'
 {'x': <class 'str'>, 'y': <class 'str'>, 'return': tuple[str, str]}
 True
 
+```
+
+## Error
+
+```text
+Runtime type is 'Node'
+
 ```