Refactored _node_replace.py InputMap/OutputMap to use a TypedDict instead of objects, simplified the schema sent to the frontend, updated nodes_post_processing.py replacements to use new schema

comfy_api/latest/_node_replace.py

@@ -1,6 +1,32 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, TypedDict
+
+
+class InputMapOldId(TypedDict):
+    """Map an old node input to a new node input by ID."""
+    new_id: str
+    old_id: str
+
+
+class InputMapSetValue(TypedDict):
+    """Set a specific value for a new node input."""
+    new_id: str
+    set_value: Any
+
+
+InputMap = InputMapOldId | InputMapSetValue
+"""
+Input mapping for node replacement. Type is inferred by dictionary keys:
+- {"new_id": str, "old_id": str} - maps old input to new input
+- {"new_id": str, "set_value": Any} - sets a specific value for new input
+"""
+
+
+class OutputMap(TypedDict):
+    """Map outputs of node replacement via indexes."""
+    new_idx: int
+    old_idx: int
 
 
 class NodeReplace:
@@ -28,67 +54,6 @@ class NodeReplace:
             "new_node_id": self.new_node_id,
             "old_node_id": self.old_node_id,
             "old_widget_ids": self.old_widget_ids,
-            "input_mapping": [m.as_dict() for m in self.input_mapping] if self.input_mapping else None,
+            "input_mapping": list(self.input_mapping) if self.input_mapping else None,
-            "output_mapping": [m.as_dict() for m in self.output_mapping] if self.output_mapping else None,
+            "output_mapping": list(self.output_mapping) if self.output_mapping else None,
-        }
-
-
-class InputMap:
-    """
-    Map inputs of node replacement.
-
-    Use InputMap.OldId or InputMap.SetValue for mapping purposes.
-    """
-    class _Assign:
-        def __init__(self, assign_type: str):
-            self.assign_type = assign_type
-
-        def as_dict(self):
-            return {
-                "assign_type": self.assign_type,
-            }
-
-    class OldId(_Assign):
-        """Connect the input of the old node with given id to new node when replacing."""
-        def __init__(self, old_id: str):
-            super().__init__("old_id")
-            self.old_id = old_id
-
-        def as_dict(self):
-            return super().as_dict() | {
-                "old_id": self.old_id,
-            }
-
-    class SetValue(_Assign):
-        """Use the given value for the input of the new node when replacing; assumes input is a widget."""
-        def __init__(self, value: Any):
-            super().__init__("set_value")
-            self.value = value
-
-        def as_dict(self):
-            return super().as_dict() | {
-                "value": self.value,
-            }
-
-    def __init__(self, new_id: str, assign: OldId | SetValue):
-        self.new_id = new_id
-        self.assign = assign
-
-    def as_dict(self):
-        return {
-            "new_id": self.new_id,
-            "assign": self.assign.as_dict(),
-        }
-
-
-class OutputMap:
-    """Map outputs of node replacement via indexes, as that's how outputs are stored."""
-    def __init__(self, new_idx: int, old_idx: int):
-        self.new_idx = new_idx
-        self.old_idx = old_idx
-
-    def as_dict(self):
-        return {
-            "new_idx": self.new_idx,
-            "old_idx": self.old_idx,
         }

comfy_extras/nodes_post_processing.py

@@ -681,12 +681,12 @@ def register_replacements_longeredge():
             old_node_id="ResizeImagesByLongerEdge",
             old_widget_ids=["longer_edge"],
             input_mapping=[
-                node_replace.InputMap(new_id="image", assign=node_replace.InputMap.OldId("images")),
+                {"new_id": "image", "old_id": "images"},
-                node_replace.InputMap(new_id="largest_size", assign=node_replace.InputMap.OldId("longer_edge")),
+                {"new_id": "largest_size", "old_id": "longer_edge"},
-                node_replace.InputMap(new_id="upscale_method", assign=node_replace.InputMap.SetValue("lanczos")),
+                {"new_id": "upscale_method", "set_value": "lanczos"},
             ],
             # just to test the frontend output_mapping code, does nothing really here
-            output_mapping=[node_replace.OutputMap(new_idx=0, old_idx=0)],
+            output_mapping=[{"new_idx": 0, "old_idx": 0}],
         ))
 
 def register_replacements_batchimages():
@@ -695,8 +695,8 @@ def register_replacements_batchimages():
             new_node_id="BatchImagesNode",
             old_node_id="ImageBatch",
             input_mapping=[
-                node_replace.InputMap(new_id="images.image0", assign=node_replace.InputMap.OldId("image1")),
+                {"new_id": "images.image0", "old_id": "image1"},
-                node_replace.InputMap(new_id="images.image1", assign=node_replace.InputMap.OldId("image2")),
+                {"new_id": "images.image1", "old_id": "image2"},
             ],
         ))
 
@@ -707,10 +707,10 @@ def register_replacements_upscaleimage():
             old_node_id="ImageScaleBy",
             old_widget_ids=["upscale_method", "scale_by"],
             input_mapping=[
-                node_replace.InputMap(new_id="input", assign=node_replace.InputMap.OldId("image")),
+                {"new_id": "input", "old_id": "image"},
-                node_replace.InputMap(new_id="resize_type", assign=node_replace.InputMap.SetValue("scale by multiplier")),
+                {"new_id": "resize_type", "set_value": "scale by multiplier"},
-                node_replace.InputMap(new_id="resize_type.multiplier", assign=node_replace.InputMap.OldId("scale_by")),
+                {"new_id": "resize_type.multiplier", "old_id": "scale_by"},
-                node_replace.InputMap(new_id="scale_method", assign=node_replace.InputMap.OldId("upscale_method")),
+                {"new_id": "scale_method", "old_id": "upscale_method"},
             ],
         ))
 
@@ -720,7 +720,7 @@ def register_replacements_controlnet():
             new_node_id="ControlNetLoader",
             old_node_id="T2IAdapterLoader",
             input_mapping=[
-                node_replace.InputMap(new_id="control_net_name", assign=node_replace.InputMap.OldId("t2i_adapter_name")),
+                {"new_id": "control_net_name", "old_id": "t2i_adapter_name"},
             ],
         ))