From 16dd7d115c8efad5f547400f989373d3335757d3 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Mon, 1 Jun 2026 18:34:30 -0700 Subject: [PATCH] DynamicSlot: address code review Bug fixes: - execution.py validate_inputs() now calls _select_option for DynamicSlot links instead of validating against slotType. Old code accepted any concrete upstream type whenever AnyType was enumerated (slotType contained '*', which validate_node_input treats as accept-anything). Verified end-to-end: LATENT into an IMAGE+AnyType slot is now rejected. - Thread live_input_types through get_input_data so custom V3 validate_inputs() sees the same DynamicSlot branch that finalization picked, instead of re-finalizing without resolver context. - DynamicSlot.Option._when_types is now an ordered tuple (preserves author declaration order); _slot_io_type/slotType ordering was previously nondeterministic via frozenset iteration. Design: - DynamicSlot.Input.get_all() now returns [self] + children, matching Autogrow / DynamicCombo so consumers like PriceBadge work uniformly. - Enforce per-option type uniqueness in DynamicSlot.Input: each io_type may appear in at most one option's 'when', and at most one option may declare when=None. Removes the ambiguous first-match-on-overlap case for single concrete types; ordering still matters when upstream is a multi-type union. - Reject non-Option entries in options=[...] explicitly. Polish: - Trim verbose DynamicSlot docstrings and inline comments. Amp-Thread-ID: https://ampcode.com/threads/T-019e8568-f382-743d-a97f-0de3ff29d501 Co-authored-by: Amp --- comfy_api/latest/_io.py | 101 ++++++++++-------- execution.py | 28 +++-- .../comfy_api_test/test_dynamic_slot.py | 93 ++++++++++++---- 3 files changed, 145 insertions(+), 77 deletions(-) diff --git a/comfy_api/latest/_io.py b/comfy_api/latest/_io.py index 6eaf7850a..a736863fd 100644 --- a/comfy_api/latest/_io.py +++ b/comfy_api/latest/_io.py @@ -1198,26 +1198,19 @@ class DynamicCombo(ComfyTypeI): class DynamicSlot(ComfyTypeI): """A slot whose revealed inputs depend on the type connected upstream. - Options dispatch on the type resolved by - :py:class:`comfy_execution.type_resolver.TypeResolver`: + Each ``Option`` declares a ``when`` condition; the first option whose + condition matches the slot's resolved upstream type (or whose + ``when=None`` matches an empty slot) decides which child inputs are + exposed. - * ``Option(when=None, ...)`` — nothing connected to the slot. - * ``Option(when=io.AnyType, ...)`` — link present, upstream resolves to ``"*"`` - (e.g. Reroute, generic If/Else, V1 nodes that declare AnyType outputs). - * ``Option(when=io.Image, ...)`` — upstream resolves to ``IMAGE``. - * ``Option(when=[io.Image, io.Mask], ...)`` — share children across types. - * ``Option(when=io.MultiType.Input(...), ...)`` — same as the list form. + Each concrete type may appear in at most one option's ``when``, so the + matching branch is unambiguous. The unconnected case (``when=None``) is + its own bucket and may also appear at most once. - On a connected slot the first option whose ``when`` set intersects the - resolved type set wins; on an unconnected slot the first ``when=None`` - option wins. No implicit "match anything I didn't enumerate" fallback — - declare ``when=io.AnyType`` if you want it. - - Known limitation: when an upstream node declares its output as ``AnyType`` - (Reroute, generic forwarders, many V1 utilities) the resolver can only - report ``"*"`` — it cannot introspect the runtime value to recover a more - specific type. Such links will always select the ``when=io.AnyType`` - branch (or no branch), never a concrete-type branch. + The AnyType limitation documented in + :py:mod:`comfy_execution.type_resolver` applies: an upstream output + declared as ``AnyType`` resolves to ``"*"`` and will only match a + ``when=io.AnyType`` option, never a concrete-type one. """ Type = dict[str, Any] @@ -1234,17 +1227,22 @@ class DynamicSlot(ComfyTypeI): def __init__(self, when: Any, inputs: list[Input]): self.when = when self.inputs = inputs + # ``_when_types`` is the ordered tuple of io_types (deterministic); + # ``_when_set`` is the same content as a set for fast matching. self._when_types = self._normalize_when(when) + self._when_set: frozenset[str] | None = ( + None if self._when_types is None else frozenset(self._when_types) + ) @staticmethod - def _normalize_when(when: Any) -> frozenset[str] | None: - """Normalize ``when`` to a ``frozenset[str]`` of io_types, or ``None`` for the unconnected case.""" + def _normalize_when(when: Any) -> tuple[str, ...] | None: + """Normalize ``when`` to an ordered, deduplicated tuple of io_types, or ``None`` for the unconnected case.""" if when is None: return None if isinstance(when, type) and issubclass(when, _ComfyType): - return frozenset([when.io_type]) + return (when.io_type,) if isinstance(when, MultiType.Input): - return frozenset(t.io_type for t in when.io_types) + return tuple(dict.fromkeys(t.io_type for t in when.io_types)) if isinstance(when, Iterable) and not isinstance(when, str): types: list[str] = [] for t in when: @@ -1252,10 +1250,11 @@ class DynamicSlot(ComfyTypeI): raise ValueError( f"DynamicSlot.Option: list entries must be ComfyType classes, got {t!r}" ) - types.append(t.io_type) + if t.io_type not in types: + types.append(t.io_type) if not types: raise ValueError("DynamicSlot.Option: when=[] is not allowed; use when=None instead") - return frozenset(types) + return tuple(types) raise ValueError( "DynamicSlot.Option: when must be None, a ComfyType class, a list of ComfyType classes, " f"or a MultiType.Input; got {when!r}" @@ -1263,7 +1262,7 @@ class DynamicSlot(ComfyTypeI): def as_dict(self): return { - "when": None if self._when_types is None else sorted(self._when_types), + "when": None if self._when_types is None else list(self._when_types), "inputs": create_input_dict_v1(self.inputs), } @@ -1272,18 +1271,34 @@ class DynamicSlot(ComfyTypeI): display_name: str=None, tooltip: str=None, lazy: bool=None, extra_dict=None): if not options: raise ValueError("DynamicSlot.Input: at least one Option is required") + for opt in options: + if not isinstance(opt, DynamicSlot.Option): + raise ValueError( + f"DynamicSlot.Input: options must be DynamicSlot.Option instances, got {opt!r}" + ) super().__init__(id, display_name, True, tooltip, lazy, extra_dict) self.options = options - # Auto-derive the slot's declared connection type as the union of - # every non-None option's `when` set. Order is preserved per option, - # then deduplicated, so authors control the displayed precedence. + # Enforce uniqueness: each io_type (and the unconnected case) may + # appear in at most one option's ``when``. Also derive the slot's + # declared connection type as the ordered union of every non-None + # option's ``when`` set so authors control displayed precedence. + seen_types: set[str] = set() + seen_none = False connected_types: list[str] = [] for opt in options: if opt._when_types is None: + if seen_none: + raise ValueError("DynamicSlot.Input: only one Option may declare when=None") + seen_none = True continue for t in opt._when_types: - if t not in connected_types: - connected_types.append(t) + if t in seen_types: + raise ValueError( + f"DynamicSlot.Input: type {t!r} appears in more than one Option's `when`; " + "each type must be claimed by exactly one option" + ) + seen_types.add(t) + connected_types.append(t) if not connected_types: raise ValueError( "DynamicSlot.Input: at least one Option must have a non-None `when`; " @@ -1291,21 +1306,19 @@ class DynamicSlot(ComfyTypeI): ) self._slot_io_type = ",".join(connected_types) - # NOTE: do NOT override get_io_type — parse_class_inputs uses the class - # io_type (COMFY_DYNAMICSLOT_V3) to dispatch into the dynamic expander. - # The auto-derived connection type is published via the `slotType` field - # in as_dict() so the frontend knows what links are accepted. + # parse_class_inputs dispatches on the class io_type (COMFY_DYNAMICSLOT_V3), + # so get_all/get_io_type must not be overridden; slotType is published via as_dict. def get_all(self) -> list[Input]: seen_ids: set[str] = set() - out: list[Input] = [] + children: list[Input] = [] for opt in self.options: for inp in opt.inputs: if inp.id in seen_ids: continue seen_ids.add(inp.id) - out.append(inp) - return out + children.append(inp) + return [self] + children def as_dict(self): return super().as_dict() | prune_dict({ @@ -1321,10 +1334,13 @@ class DynamicSlot(ComfyTypeI): @staticmethod def _select_option(options: list[dict[str, Any]], live_input_types: dict[str, str] | None, finalized_id: str, has_link: bool) -> dict[str, Any] | None: - """Pick the first option whose ``when`` matches the slot's current state. + """Pick the first option whose ``when`` matches the slot's state. - Matching is set intersection against the resolved type string split on - commas (so MultiType outputs like ``"IMAGE,MASK"`` work naturally). + Connected: pick the first option whose ``when`` set intersects the + comma-split resolved type. Unconnected: pick the first ``when=None``. + With per-option type uniqueness, at most one connected option can match + any single concrete type; ordering only matters when upstream declares + a multi-type union (e.g. ``"IMAGE,MASK"``). """ if not has_link: for opt in options: @@ -1332,7 +1348,7 @@ class DynamicSlot(ComfyTypeI): return opt return None resolved = (live_input_types or {}).get(finalized_id, "*") - resolved_set = set(t.strip() for t in resolved.split(",")) + resolved_set = {t.strip() for t in resolved.split(",")} for opt in options: when = opt["when"] if when is None: @@ -1350,8 +1366,7 @@ class DynamicSlot(ComfyTypeI): if selected is not None: parse_class_inputs(out_dict, live_inputs, selected["inputs"], curr_prefix, live_input_types) # Always advertise the slot itself so the connector renders even when no - # option matched (e.g. resolved type wasn't enumerated and there's no - # AnyType option). Unmatched cases just expand no children. + # option matched (unmatched concrete + no AnyType option). out_dict[input_type][finalized_id] = value out_dict["dynamic_paths"][finalized_id] = finalize_prefix(curr_prefix, curr_prefix[-1]) diff --git a/execution.py b/execution.py index 93e0f185e..b8727473b 100644 --- a/execution.py +++ b/execution.py @@ -152,15 +152,14 @@ class CacheSet: SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org") -def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt=None, extra_data={}): +def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt=None, extra_data={}, live_input_types=None): is_v3 = issubclass(class_def, _ComfyNodeInternal) v3_data: io.V3Data = {} hidden_inputs_v3 = {} valid_inputs = class_def.INPUT_TYPES() if is_v3: # Let dynamic schemas branch on resolved upstream types, not just literal values. - live_input_types = None - if dynprompt is not None and hasattr(dynprompt, "get_type_resolver"): + if live_input_types is None and dynprompt is not None and hasattr(dynprompt, "get_type_resolver"): live_input_types = dynprompt.get_type_resolver().compute_live_input_types(unique_id) valid_inputs, hidden, v3_data = _io.get_finalized_class_inputs(valid_inputs, inputs, live_input_types=live_input_types) input_data_all = {} @@ -867,6 +866,7 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None, typ v3_data = None validate_function_inputs = [] validate_has_kwargs = False + live_input_types = None if issubclass(obj_class, _ComfyNodeInternal): obj_class: _io._ComfyNodeBaseInternal class_inputs = obj_class.INPUT_TYPES() @@ -926,14 +926,16 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None, typ # frontend-injected type metadata get the same answer as the UI. received_type = type_resolver.resolve_output_type(o_id, val[1]) received_types[x] = received_type - # DynamicSlot publishes its accepted connection types via the - # `slotType` field (auto-derived union of option `when` types). - # The declared input_type ("COMFY_DYNAMICSLOT_V3") is just the - # dispatch tag; validate against slotType instead. - effective_input_type = input_type + # DynamicSlot's declared input_type is just the dispatch tag + # (COMFY_DYNAMICSLOT_V3); a link is valid iff some Option would + # actually claim the resolved upstream type. if input_type == _io.DynamicSlot.io_type and isinstance(extra_info, dict): - effective_input_type = extra_info.get("slotType", input_type) - if 'input_types' not in validate_function_inputs and not validate_node_input(received_type, effective_input_type): + link_valid = _io.DynamicSlot._select_option( + extra_info.get("options", []), {x: received_type}, x, has_link=True + ) is not None + else: + link_valid = validate_node_input(received_type, input_type) + if 'input_types' not in validate_function_inputs and not link_valid: details = f"{x}, received_type({received_type}) mismatch input_type({input_type})" error = { "type": "return_type_mismatch", @@ -1079,7 +1081,11 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None, typ continue if len(validate_function_inputs) > 0 or validate_has_kwargs: - input_data_all, _, v3_data = get_input_data(inputs, obj_class, unique_id) + # Reuse the precomputed live_input_types so a custom validate_inputs() + # sees the same DynamicSlot branch that finalization picked above. + input_data_all, _, v3_data = get_input_data( + inputs, obj_class, unique_id, live_input_types=live_input_types + ) input_filtered = {} for x in input_data_all: if x in validate_function_inputs or validate_has_kwargs: diff --git a/tests-unit/comfy_api_test/test_dynamic_slot.py b/tests-unit/comfy_api_test/test_dynamic_slot.py index ef0f520c2..c4444d745 100644 --- a/tests-unit/comfy_api_test/test_dynamic_slot.py +++ b/tests-unit/comfy_api_test/test_dynamic_slot.py @@ -19,32 +19,39 @@ def _opt(when, ids=None): def test_option_when_none(): o = _opt(None, ["a"]) assert o._when_types is None + assert o._when_set is None assert o.as_dict()["when"] is None def test_option_when_single_type(): o = _opt(io.Image) - assert o._when_types == frozenset({"IMAGE"}) + assert o._when_types == ("IMAGE",) + assert o._when_set == frozenset({"IMAGE"}) assert o.as_dict()["when"] == ["IMAGE"] def test_option_when_anytype(): o = _opt(io.AnyType) - assert o._when_types == frozenset({"*"}) + assert o._when_types == ("*",) assert o.as_dict()["when"] == ["*"] -def test_option_when_list(): - o = _opt([io.Image, io.Mask]) - assert o._when_types == frozenset({"IMAGE", "MASK"}) - # list form sorted for stable serialization - assert o.as_dict()["when"] == ["IMAGE", "MASK"] +def test_option_when_list_preserves_order(): + """Declaration order is preserved in both the tuple and the serialized form.""" + o = _opt([io.Mask, io.Image]) + assert o._when_types == ("MASK", "IMAGE") + assert o.as_dict()["when"] == ["MASK", "IMAGE"] + + +def test_option_when_list_dedups_within_option(): + o = _opt([io.Image, io.Image, io.Mask]) + assert o._when_types == ("IMAGE", "MASK") def test_option_when_multitype_input(): mt = io.MultiType.Input("x", types=[io.Image, io.Latent]) o = _opt(mt) - assert o._when_types == frozenset({"IMAGE", "LATENT"}) + assert o._when_types == ("IMAGE", "LATENT") def test_option_when_empty_list_rejected(): @@ -83,10 +90,6 @@ def test_input_auto_derives_slot_type(): _opt(None, ["c"]), ]) # Declared order preserved across non-None options; None contributes nothing. - # Note: get_io_type() intentionally still returns the dynamic class io_type - # (COMFY_DYNAMICSLOT_V3) so parse_class_inputs dispatches into the expander. - # The auto-derived slot type is exposed via the `slotType` field of as_dict() - # and via the private `_slot_io_type` attribute (used by the type resolver). assert inp._slot_io_type == "IMAGE,MASK" d = inp.as_dict() assert d["slotType"] == "IMAGE,MASK" @@ -101,13 +104,45 @@ def test_input_includes_anytype_in_slot_type(): assert inp._slot_io_type == "IMAGE,*" -def test_input_get_all_dedups_inputs_by_id(): +def test_input_rejects_duplicate_type_across_options(): + with pytest.raises(ValueError, match="appears in more than one"): + io.DynamicSlot.Input("x", options=[ + _opt(io.Image, ["a"]), + _opt([io.Image, io.Mask], ["b"]), + ]) + + +def test_input_rejects_duplicate_anytype_across_options(): + with pytest.raises(ValueError, match="appears in more than one"): + io.DynamicSlot.Input("x", options=[ + _opt(io.AnyType, ["a"]), + _opt(io.AnyType, ["b"]), + ]) + + +def test_input_rejects_duplicate_when_none(): + with pytest.raises(ValueError, match="only one Option may declare when=None"): + io.DynamicSlot.Input("x", options=[ + _opt(io.Image, ["a"]), + _opt(None, ["b"]), + _opt(None, ["c"]), + ]) + + +def test_input_rejects_non_option_entry(): + with pytest.raises(ValueError, match="must be DynamicSlot.Option instances"): + io.DynamicSlot.Input("x", options=[_opt(io.Image, ["a"]), "not an option"]) + + +def test_input_get_all_prepends_self_and_dedups_children(): inp = io.DynamicSlot.Input("x", options=[ _opt(io.Image, ["shared", "image_only"]), _opt(io.Mask, ["shared", "mask_only"]), ]) - ids = [i.id for i in inp.get_all()] - assert ids == ["shared", "image_only", "mask_only"] + items = inp.get_all() + # Convention shared with Autogrow / DynamicCombo: parent first, then children. + assert items[0] is inp + assert [i.id for i in items[1:]] == ["shared", "image_only", "mask_only"] # --------------------------------------------------------------------------- @@ -156,17 +191,29 @@ def test_select_anytype_does_not_match_concrete(): assert _select(options, {"x": "MASK"}, has_link=True) is None -def test_select_first_match_wins(): +def test_select_anytype_branch_does_not_swallow_unenumerated_concrete(): + """Regression: a slot exposing IMAGE + AnyType must reject LATENT upstream + instead of expanding the AnyType branch. validate_inputs relies on this + to compute link validity (slotType="IMAGE,*" alone would over-accept).""" + options = [_opt(io.Image, ["image_widget"]), _opt(io.AnyType, ["any_widget"])] + assert _select(options, {"x": "LATENT"}, has_link=True) is None + # Sanity: IMAGE still matches the IMAGE branch and "*" still matches AnyType. + assert _select(options, {"x": "IMAGE"}, has_link=True)["when"] == ["IMAGE"] + assert _select(options, {"x": "*"}, has_link=True)["when"] == ["*"] + + +def test_select_first_match_wins_on_union_upstream(): + """Ordering only matters when upstream declares a multi-type union; with + per-option type uniqueness, single concrete types can never match two + options.""" options = [ - _opt([io.Image, io.Mask], ["both"]), - _opt(io.Image, ["image_only"]), + _opt([io.Image, io.Mask], ["image_or_mask"]), + _opt(io.Latent, ["latent_only"]), ] - # Resolved IMAGE matches both; first option wins. - sel = _select(options, {"x": "IMAGE"}, has_link=True) - assert sel["inputs"] - # The "both" option's first input is named "both" + # Upstream union "IMAGE,LATENT" intersects both options; first option wins. + sel = _select(options, {"x": "IMAGE,LATENT"}, has_link=True) first_input_id = next(iter(sel["inputs"]["required"].keys())) - assert first_input_id == "both" + assert first_input_id == "image_or_mask" def test_select_multitype_upstream_intersects_option_set():