Skip to content

Commit

Permalink
feat(nodes): improve types in graph.py
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
psychedelicious committed Feb 13, 2024
1 parent 8bd65be commit 3726293
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions invokeai/app/services/shared/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -540,7 +540,7 @@ def has_node(self, node_path: str) -> bool:
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)
Expand Down Expand Up @@ -891,7 +891,7 @@ def next(self) -> Optional[BaseInvocation]:
# 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:
Expand Down

0 comments on commit 3726293

Please sign in to comment.