clean up duplicate code from merge conflict

Author: jab416171

src/comfy_sdk/workflows.py

@@ -174,104 +174,6 @@ def add_node(
 
         return node_id
 
-    def remove_node(self, node_id: str) -> None:
-        """Remove a node and redirect links through it back to their sources.
-
-        Deletes the node identified by ``node_id`` from the graph. Any input
-        connections (links) in other nodes that reference this node's outputs
-        are redirected to the source that fed into the removed node, effectively
-        unwinding any insertion point.
-
-        If the removed node has exactly one input that is a link, all downstream
-        consumers of its outputs are redirected to that source. Otherwise (zero
-        or multiple link inputs), downstream links are simply deleted.
-        """
-        removed = self.json.pop(node_id, None)
-        if removed is None:
-            return
-
-        # Collect link inputs from the removed node
-        link_inputs: list[tuple[str, int]] = []
-        removed_inputs = removed.get("inputs") or {}
-        for value in removed_inputs.values():
-            if _is_link(value):
-                link_inputs.append((value[0], int(value[1])))
-
-        if len(link_inputs) == 1:
-            # Single link input: redirect all downstream consumers to that source
-            src_node, src_output = link_inputs[0]
-            for node in self.json.values():
-                inputs = node.get("inputs")
-                if not inputs:
-                    continue
-                for key, value in list(inputs.items()):
-                    if _is_link(value) and value[0] == node_id:
-                        if src_node in self.json:
-                            inputs[key] = [src_node, src_output]
-                        else:
-                            del inputs[key]
-        else:
-            # Zero or multiple link inputs: just delete downstream links
-            for node in self.json.values():
-                inputs = node.get("inputs")
-                if not inputs:
-                    continue
-                to_delete = []
-                for key, value in inputs.items():
-                    if _is_link(value) and value[0] == node_id:
-                        to_delete.append(key)
-                for key in to_delete:
-                    del inputs[key]
-
-    def remove_node(self, node_id: str) -> None:
-        """Remove a node and redirect links through it back to their sources.
-
-        Deletes the node identified by ``node_id`` from the graph. Any input
-        connections (links) in other nodes that reference this node's outputs
-        are redirected to the source that fed into the removed node, effectively
-        unwinding any insertion point.
-
-        If the removed node has exactly one input that is a link, all downstream
-        consumers of its outputs are redirected to that source. Otherwise (zero
-        or multiple link inputs), downstream links are simply deleted.
-        """
-        removed = self.json.pop(node_id, None)
-        if removed is None:
-            return
-
-        # Collect link inputs from the removed node
-        link_inputs: list[tuple[str, int]] = []
-        removed_inputs = removed.get("inputs") or {}
-        for value in removed_inputs.values():
-            if _is_link(value):
-                link_inputs.append((value[0], int(value[1])))
-
-        if len(link_inputs) == 1:
-            # Single link input: redirect all downstream consumers to that source
-            src_node, src_output = link_inputs[0]
-            for node in self.json.values():
-                inputs = node.get("inputs")
-                if not inputs:
-                    continue
-                for key, value in list(inputs.items()):
-                    if _is_link(value) and value[0] == node_id:
-                        if src_node in self.json:
-                            inputs[key] = [src_node, src_output]
-                        else:
-                            del inputs[key]
-        else:
-            # Zero or multiple link inputs: just delete downstream links
-            for node in self.json.values():
-                inputs = node.get("inputs")
-                if not inputs:
-                    continue
-                to_delete = []
-                for key, value in inputs.items():
-                    if _is_link(value) and value[0] == node_id:
-                        to_delete.append(key)
-                for key in to_delete:
-                    del inputs[key]
-
     def __repr__(self) -> str:
         return f"Workflow(nodes={len(self.json)})"
 

tests/test_workflows.py

@@ -199,136 +199,6 @@ def test_add_node_auto_id_on_empty_graph():
     assert new_id == "1"
 
 
-def test_add_node_redirects_downstream_single_consumer():
-    graph = {
-        "1": {
-            "class_type": "UNETLoader",
-            "inputs": {"unet_name": "model.safetensors"},
-        },
-        "2": {
-            "class_type": "LoraLoaderModelOnly",
-            "inputs": {
-                "lora_name": "lora.safetensors",
-                "strength_model": 1,
-                "model": ["1", 0],
-            },
-        },
-        "4": {
-            "class_type": "KSampler",
-            "inputs": {
-                "seed": 0,
-                "model": ["2", 0],
-            },
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "ModelAttentionBackend",
-        before="4",
-        inputs={
-            "attention": "pytorch attention",
-            "model": ["2", 0],
-        },
-    )
-    assert new_id == "5"
-    assert wf.json[new_id]["class_type"] == "ModelAttentionBackend"
-    assert wf.json[new_id]["inputs"]["model"] == ["2", 0]
-    assert wf.json["4"]["inputs"]["model"] == [new_id, 0]
-
-
-def test_add_node_redirects_multiple_downstream():
-    graph = {
-        "1": {
-            "class_type": "LoadImage",
-            "inputs": {"image": "example.png"},
-        },
-        "2": {
-            "class_type": "PreviewImage",
-            "inputs": {"images": ["1", 0]},
-        },
-        "3": {
-            "class_type": "PreviewImage",
-            "inputs": {"images": ["1", 0]},
-        },
-        "4": {
-            "class_type": "PreviewImage",
-            "inputs": {"images": ["1", 0]},
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "ImageScaleToTotalPixels",
-        after="1",
-        inputs={
-            "upscale_method": "nearest-exact",
-            "megapixels": 1,
-            "image": ["1", 0],
-        },
-    )
-    assert new_id == "5"
-    assert wf.json[new_id]["inputs"]["image"] == ["1", 0]
-    assert wf.json["2"]["inputs"]["images"] == [new_id, 0]
-    assert wf.json["3"]["inputs"]["images"] == [new_id, 0]
-    assert wf.json["4"]["inputs"]["images"] == [new_id, 0]
-
-
-def test_add_node_no_redirect_when_upstream_not_in_graph():
-    graph = {
-        "1": {
-            "class_type": "CLIPTextEncode",
-            "inputs": {"text": "hello"},
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "KSampler",
-        inputs={"model": ["999", 0]},
-    )
-    assert new_id == "2"
-    assert wf.json[new_id]["inputs"]["model"] == ["999", 0]
-
-
-def test_add_node_no_redirect_when_no_downstream_consumers():
-    graph = {
-        "1": {
-            "class_type": "CheckpointLoader",
-            "inputs": {"ckpt_name": "model.safetensors"},
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "LoraLoader",
-        inputs={"model": ["1", 0], "clip": ["1", 1]},
-    )
-    assert new_id == "2"
-    assert wf.json[new_id]["inputs"]["model"] == ["1", 0]
-    assert wf.json[new_id]["inputs"]["clip"] == ["1", 1]
-
-
-def test_add_node_both_before_and_after_raises():
-    wf = Workflow({"1": {"class_type": "X", "inputs": {}}})
-    try:
-        wf.add_node("Y", before="1", after="1", inputs={})
-    except ValueError as e:
-        assert "not both" in str(e)
-    else:
-        assert False, "Expected ValueError"
-
-
-def test_add_node_no_inputs_no_redirect():
-    wf = Workflow({"1": {"class_type": "X", "inputs": {}}})
-    new_id = wf.add_node("Y")
-    assert new_id == "2"
-    assert wf.json[new_id]["class_type"] == "Y"
-    assert "inputs" not in wf.json[new_id]
-
-
-def test_add_node_auto_id_on_empty_graph():
-    wf = Workflow({})
-    new_id = wf.add_node("X")
-    assert new_id == "1"
-
-
 def test_remove_node_model_attention_backend():
     graph = {
         "1": {
@@ -469,133 +339,3 @@ def test_remove_node_redirects_preview_any():
 
     assert "2" not in wf.json
     assert wf.json["3"]["inputs"]["text"] == ["1", 0]
-
-
-def test_add_node_redirects_downstream_single_consumer():
-    graph = {
-        "1": {
-            "class_type": "UNETLoader",
-            "inputs": {"unet_name": "model.safetensors"},
-        },
-        "2": {
-            "class_type": "LoraLoaderModelOnly",
-            "inputs": {
-                "lora_name": "lora.safetensors",
-                "strength_model": 1,
-                "model": ["1", 0],
-            },
-        },
-        "4": {
-            "class_type": "KSampler",
-            "inputs": {
-                "seed": 0,
-                "model": ["2", 0],
-            },
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "ModelAttentionBackend",
-        before="4",
-        inputs={
-            "attention": "pytorch attention",
-            "model": ["2", 0],
-        },
-    )
-    assert new_id == "5"
-    assert wf.json[new_id]["class_type"] == "ModelAttentionBackend"
-    assert wf.json[new_id]["inputs"]["model"] == ["2", 0]
-    assert wf.json["4"]["inputs"]["model"] == [new_id, 0]
-
-
-def test_add_node_redirects_multiple_downstream():
-    graph = {
-        "1": {
-            "class_type": "LoadImage",
-            "inputs": {"image": "example.png"},
-        },
-        "2": {
-            "class_type": "PreviewImage",
-            "inputs": {"images": ["1", 0]},
-        },
-        "3": {
-            "class_type": "PreviewImage",
-            "inputs": {"images": ["1", 0]},
-        },
-        "4": {
-            "class_type": "PreviewImage",
-            "inputs": {"images": ["1", 0]},
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "ImageScaleToTotalPixels",
-        after="1",
-        inputs={
-            "upscale_method": "nearest-exact",
-            "megapixels": 1,
-            "image": ["1", 0],
-        },
-    )
-    assert new_id == "5"
-    assert wf.json[new_id]["inputs"]["image"] == ["1", 0]
-    assert wf.json["2"]["inputs"]["images"] == [new_id, 0]
-    assert wf.json["3"]["inputs"]["images"] == [new_id, 0]
-    assert wf.json["4"]["inputs"]["images"] == [new_id, 0]
-
-
-def test_add_node_no_redirect_when_upstream_not_in_graph():
-    graph = {
-        "1": {
-            "class_type": "CLIPTextEncode",
-            "inputs": {"text": "hello"},
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "KSampler",
-        inputs={"model": ["999", 0]},
-    )
-    assert new_id == "2"
-    assert wf.json[new_id]["inputs"]["model"] == ["999", 0]
-
-
-def test_add_node_no_redirect_when_no_downstream_consumers():
-    graph = {
-        "1": {
-            "class_type": "CheckpointLoader",
-            "inputs": {"ckpt_name": "model.safetensors"},
-        },
-    }
-    wf = Workflow(graph)
-    new_id = wf.add_node(
-        "LoraLoader",
-        inputs={"model": ["1", 0], "clip": ["1", 1]},
-    )
-    assert new_id == "2"
-    assert wf.json[new_id]["inputs"]["model"] == ["1", 0]
-    assert wf.json[new_id]["inputs"]["clip"] == ["1", 1]
-
-
-def test_add_node_both_before_and_after_raises():
-    wf = Workflow({"1": {"class_type": "X", "inputs": {}}})
-    try:
-        wf.add_node("Y", before="1", after="1", inputs={})
-    except ValueError as e:
-        assert "not both" in str(e)
-    else:
-        assert False, "Expected ValueError"
-
-
-def test_add_node_no_inputs_no_redirect():
-    wf = Workflow({"1": {"class_type": "X", "inputs": {}}})
-    new_id = wf.add_node("Y")
-    assert new_id == "2"
-    assert wf.json[new_id]["class_type"] == "Y"
-    assert "inputs" not in wf.json[new_id]
-
-
-def test_add_node_auto_id_on_empty_graph():
-    wf = Workflow({})
-    new_id = wf.add_node("X")
-    assert new_id == "1"