feat(nodes): improve types in graph.py

Methods `get_node` and `complete` were typed as returning a dynamically created unions `InvocationsUnion` and `InvocationOutputsUnion`, respectively.

Static type analysers cannot work with dynamic objects, so these methods end up as effectively un-annotated, returning `Unknown`.

They now return `BaseInvocation` and `BaseInvocationOutput`, respectively, which are the superclasses of all members of each union. This gives us the best type annotation that is possible.

Note: the return types of these methods are never introspected, so it doesn't really matter what they are at runtime.
This commit is contained in:
psychedelicious 2024-02-05 18:27:31 +11:00
parent 8bd65be8c8
commit 3726293258

View File

@ -540,7 +540,7 @@ class Graph(BaseModel):
except NodeNotFoundError:
return False
def get_node(self, node_path: str) -> InvocationsUnion:
def get_node(self, node_path: str) -> BaseInvocation:
"""Gets a node from the graph using a node path."""
# Materialized graphs may have nodes at the top level
graph, node_id = self._get_graph_and_node(node_path)
@ -891,7 +891,7 @@ class GraphExecutionState(BaseModel):
# If next is still none, there's no next node, return None
return next_node
def complete(self, node_id: str, output: InvocationOutputsUnion):
def complete(self, node_id: str, output: BaseInvocationOutput) -> None:
"""Marks a node as complete"""
if node_id not in self.execution_graph.nodes: