Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Return both nodes on replacement #560

Merged
merged 2 commits into from
Jan 17, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 23 additions & 11 deletions pyiron_workflow/nodes/composite.py
Original file line number Diff line number Diff line change
Expand Up @@ -326,7 +326,7 @@ def remove_child(self, child: Node | str) -> Node:

def replace_child(
self, owned_node: Node | str, replacement: Node | type[Node]
) -> Node:
) -> tuple[Node, Node]:
"""
Replaces a node currently owned with a new node instance.
The replacement must not belong to any other parent or have any connections.
Expand All @@ -348,7 +348,7 @@ def replace_child(
and simply gets instantiated.)

Returns:
(Node): The node that got removed
(Node, Node): The node that got removed and the new one that replaced it.
"""
if isinstance(owned_node, str):
owned_node = self.children[owned_node]
Expand All @@ -367,44 +367,56 @@ def replace_child(
)
if replacement.connected:
raise ValueError("Replacement node must not have any connections")
replacement_node = replacement
elif issubclass(replacement, Node):
replacement = replacement(label=owned_node.label)
replacement_node = replacement(label=owned_node.label)
else:
raise TypeError(
f"Expected replacement node to be a node instance or node subclass, but "
f"got {replacement}"
)

replacement.copy_io(owned_node) # If the replacement is incompatible, we'll
replacement_node.copy_io(
owned_node
) # If the replacement is incompatible, we'll
# fail here before we've changed the parent at all. Since the replacement was
# first guaranteed to be an unconnected orphan, there is not yet any permanent
# damage
is_starting_node = owned_node in self.starting_nodes
# In case the replaced node interfaces with the composite's IO, catch value
# links
inbound_links = [
(sending_channel, replacement.inputs[sending_channel.value_receiver.label])
(
sending_channel,
replacement_node.inputs[sending_channel.value_receiver.label],
)
for sending_channel in self.inputs
if sending_channel.value_receiver in owned_node.inputs
]
outbound_links = [
(replacement.outputs[sending_channel.label], sending_channel.value_receiver)
(
replacement_node.outputs[sending_channel.label],
sending_channel.value_receiver,
)
for sending_channel in owned_node.outputs
if sending_channel.value_receiver in self.outputs
]
self.remove_child(owned_node)
replacement.label, owned_node.label = owned_node.label, replacement.label
self.add_child(replacement)
replacement_node.label, owned_node.label = (
owned_node.label,
replacement_node.label,
)
self.add_child(replacement_node)
if is_starting_node:
self.starting_nodes.append(replacement)
self.starting_nodes.append(replacement_node)
for sending_channel, receiving_channel in inbound_links + outbound_links:
sending_channel.value_receiver = receiving_channel

# Clear caches
self._cached_inputs = None
replacement._cached_inputs = None
replacement_node._cached_inputs = None

return owned_node
return owned_node, replacement_node

def executor_shutdown(self, wait=True, *, cancel_futures=False):
"""
Expand Down
2 changes: 1 addition & 1 deletion pyiron_workflow/nodes/macro.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,7 @@ class Macro(Composite, StaticNode, ScrapesIO, ABC):
>>> # With the replace method
>>> # (replacement target can be specified by label or instance,
>>> # the replacing node can be specified by instance or class)
>>> replaced = adds_six_macro.replace_child(adds_six_macro.one, add_two())
>>> replaced, _ = adds_six_macro.replace_child(adds_six_macro.one, add_two())
>>> # With the replace_with method
>>> adds_six_macro.two.replace_with(add_two())
>>> # And by assignment of a compatible class to an occupied node label
Expand Down
10 changes: 6 additions & 4 deletions pyiron_workflow/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -497,8 +497,10 @@ def _owned_io_panels(self) -> list[IO]:

def replace_child(
self, owned_node: Node | str, replacement: Node | type[Node]
) -> Node:
super().replace_child(owned_node=owned_node, replacement=replacement)
) -> tuple[Node, Node]:
replaced, replacement_node = super().replace_child(
owned_node=owned_node, replacement=replacement
)

# Finally, make sure the IO is constructible with this new node, which will
# catch things like incompatible IO maps
Expand All @@ -509,11 +511,11 @@ def replace_child(
except Exception as e:
# If IO can't be successfully rebuilt using this node, revert changes and
# raise the exception
self.replace_child(replacement, owned_node) # Guaranteed to work since
self.replace_child(replacement_node, replaced) # Guaranteed to work since
# replacement in the other direction was already a success
raise e

return owned_node
return replaced, replacement_node

@property
def parent(self) -> None:
Expand Down
Loading