From 0d1b5bc133a7a553cfe4334e57171aecd6bc9732 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 19 Aug 2024 14:07:05 +0200 Subject: [PATCH] Tree.replace_node: tweak interface --- loopy/schedule/tree.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 5b830494a..ef9efd47f 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -154,15 +154,15 @@ def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": .set(node, ())), self._child_to_parent.set(node, parent)) - def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": + def replace_node(self, node: NodeT, new_node: NodeT) -> "Tree[NodeT]": """ - Returns a copy of *self* with *node* replaced with *new_id*. + Returns a copy of *self* with *node* replaced with *new_node*. """ - if not self.is_a_node(node): + if node not in self: raise ValueError(f"'{node}' not present in tree.") - if self.is_a_node(new_id): - raise ValueError(f"cannot rename to '{new_id}', as its already a part" + if new_node in self: + raise ValueError(f"cannot replace with '{new_node}', as its already a part" " of the tree.") parent = self.parent(node)