typing

Author: justinchuby

onnxscript/ir/_core.py

@@ -2718,11 +2718,11 @@ def remove(self, nodes: Node | Iterable[Node], /, safe: bool = False) -> None:
         """
         self._graph.remove(nodes, safe=safe)
 
-    def insert_after(self, node: Node, new_nodes: Iterable[Node], /) -> None:
+    def insert_after(self, node: Node, new_nodes: Iterable[Node] | Node, /) -> None:
         """Insert new nodes after the given node in O(#new_nodes) time."""
         self._graph.insert_after(node, new_nodes)
 
-    def insert_before(self, node: Node, new_nodes: Iterable[Node], /) -> None:
+    def insert_before(self, node: Node, new_nodes: Iterable[Node] | Node, /) -> None:
         """Insert new nodes before the given node in O(#new_nodes) time."""
         self._graph.insert_before(node, new_nodes)
 

onnxscript/ir/_protocols.py

@@ -320,11 +320,11 @@ def remove(self, node: NodeProtocol, /) -> None:
         """Remove a node from the graph."""
         ...
 
-    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /) -> None:
         """Insert new nodes after the given node."""
         ...
 
-    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /) -> None:
         """Insert new nodes before the given node."""
         ...
 
@@ -589,11 +589,11 @@ def remove(self, node: NodeProtocol, /) -> None:
         """Remove a node from the function."""
         ...
 
-    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_after(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /) -> None:
         """Insert new nodes after the given node."""
         ...
 
-    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol], /) -> None:
+    def insert_before(self, node: NodeProtocol, new_nodes: Iterator[NodeProtocol] | NodeProtocol, /) -> None:
         """Insert new nodes before the given node."""
         ...