@@ -134,7 +134,7 @@ def _replace_node(self, original_node: EntityType, new_node: EntityType):
134134 def _replace_subgraph (
135135 self ,
136136 graph : Optional [EntityGraph ],
137- removed_nodes : Optional [Set [EntityType ]],
137+ nodes_to_remove : Optional [Set [EntityType ]],
138138 new_results : Optional [List [Entity ]] = None ,
139139 ):
140140 """
@@ -146,7 +146,7 @@ def _replace_subgraph(
146146 ----------
147147 graph : EntityGraph, optional
148148 The input graph. If it's none, no new node and edge will be added.
149- removed_nodes : Set[EntityType], optional
149+ nodes_to_remove : Set[EntityType], optional
150150 The nodes to be removed. All the edges connected with them are removed as well.
151151 new_results : List[EntityType], optional, default None
152152 The updated results of the graph. If it's None, then the results will not be updated.
@@ -160,18 +160,18 @@ def _replace_subgraph(
160160 affected_successors = set ()
161161
162162 output_to_node = dict ()
163- removed_nodes = removed_nodes or set ()
163+ nodes_to_remove = nodes_to_remove or set ()
164164 if graph is not None :
165165 # Add the output key -> node of the subgraph
166166 for node in graph .iter_nodes ():
167- if node in removed_nodes :
167+ if node in nodes_to_remove :
168168 raise ValueError (f"The node { node } is in the removed set" )
169169 for output in node .outputs :
170170 output_to_node [output .key ] = node
171171
172- for node in removed_nodes :
172+ for node in nodes_to_remove :
173173 for affected_successor in self ._graph .iter_successors (node ):
174- if affected_successor not in removed_nodes :
174+ if affected_successor not in nodes_to_remove :
175175 affected_successors .add (affected_successor )
176176 # Check whether affected successors' inputs are in subgraph
177177 for affected_successor in affected_successors :
@@ -180,7 +180,7 @@ def _replace_subgraph(
180180 raise ValueError (
181181 f"The output { inp } of node { affected_successor } is missing in the subgraph"
182182 )
183- for node in removed_nodes :
183+ for node in nodes_to_remove :
184184 self ._graph .remove_node (node )
185185
186186 if graph is None :
@@ -200,7 +200,7 @@ def _replace_subgraph(
200200 self ._graph .add_edge (pred_node , node )
201201
202202 if new_results is not None :
203- self ._graph .results = new_results . copy ( )
203+ self ._graph .results = list ( new_results )
204204
205205 def _add_collapsable_predecessor (self , node : EntityType , predecessor : EntityType ):
206206 pred_original = self ._records .get_original_entity (predecessor , predecessor )
0 commit comments