From 3726293258b5c4e062e3c2bef9d2e0304a3ffdae Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 5 Feb 2024 18:27:31 +1100 Subject: [PATCH] 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. --- invokeai/app/services/shared/graph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 80f56b49d3..1acf165aba 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -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: