95 lines
3.3 KiB
Python
95 lines
3.3 KiB
Python
from .any_type import ANY_TYPE
|
|
|
|
class FriendlyPipeEdit:
|
|
"""
|
|
A pipe editing node that adds new slots to an existing FRIENDLY_PIPE.
|
|
Allows extending pipes with additional inputs while preserving existing data.
|
|
"""
|
|
|
|
MAX_SLOTS = 80
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"pipe": ("FRIENDLY_PIPE",),
|
|
},
|
|
"optional": {},
|
|
"hidden": {
|
|
"slot_count": ("INT", {"default": 1}),
|
|
"slot_names": ("STRING", {"default": "{}"}),
|
|
}
|
|
}
|
|
|
|
# Define additional slots as optional with ANY_TYPE
|
|
for i in range(1, cls.MAX_SLOTS + 1):
|
|
inputs["optional"][f"slot_{i}"] = (ANY_TYPE, {"forceInput": True})
|
|
|
|
return inputs
|
|
|
|
RETURN_TYPES = ("FRIENDLY_PIPE",)
|
|
RETURN_NAMES = ("pipe",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "utils/pipe"
|
|
|
|
def execute(self, pipe, slot_count=1, slot_names="{}", **kwargs):
|
|
import json
|
|
|
|
# Parse slot names from JSON string
|
|
try:
|
|
names_dict = json.loads(slot_names)
|
|
except:
|
|
names_dict = {}
|
|
|
|
# Start with the incoming pipe data
|
|
incoming_slots = pipe.get("slots", {})
|
|
incoming_names = pipe.get("names", {})
|
|
|
|
# Calculate actual incoming slot count from the data (more reliable than slot_count field)
|
|
incoming_slot_count = 0
|
|
for key in incoming_slots.keys():
|
|
int_key = int(key) if isinstance(key, str) else key
|
|
incoming_slot_count = max(incoming_slot_count, int_key)
|
|
|
|
# Also consider the pipe's reported slot_count (for empty slots)
|
|
incoming_slot_count = max(incoming_slot_count, pipe.get("slot_count", 0))
|
|
|
|
# Create new pipe data combining incoming + new slots
|
|
pipe_data = {
|
|
"slots": {},
|
|
"names": {},
|
|
}
|
|
|
|
# Copy incoming slots with integer keys (normalize from string if needed)
|
|
for key, value in incoming_slots.items():
|
|
int_key = int(key) if isinstance(key, str) else key
|
|
pipe_data["slots"][int_key] = value
|
|
|
|
# Copy incoming names with integer keys (normalize from string if needed)
|
|
for key, value in incoming_names.items():
|
|
int_key = int(key) if isinstance(key, str) else key
|
|
pipe_data["names"][int_key] = value
|
|
|
|
# Find the highest slot we're adding
|
|
max_new_slot = 0
|
|
|
|
# Add new slots with offset indices
|
|
for i in range(1, slot_count + 1):
|
|
slot_key = f"slot_{i}"
|
|
new_index = incoming_slot_count + i
|
|
|
|
if slot_key in kwargs and kwargs[slot_key] is not None:
|
|
pipe_data["slots"][new_index] = kwargs[slot_key]
|
|
max_new_slot = max(max_new_slot, i)
|
|
|
|
# Add names with offset
|
|
if i in names_dict:
|
|
pipe_data["names"][new_index] = names_dict[i]
|
|
elif str(i) in names_dict:
|
|
pipe_data["names"][new_index] = names_dict[str(i)]
|
|
|
|
# Calculate final slot count
|
|
pipe_data["slot_count"] = incoming_slot_count + max(slot_count, max_new_slot)
|
|
|
|
return (pipe_data,)
|