Pin enum serialize behavior; guard key collisions; document set limitation

The read side does not round-trip Enums: node-graph's coerce_inputs_from_spec
records structured_type extras only for dataclass/pydantic/TypedDict, never
for Enum sockets, so a task body declaring a plain Enum input receives the
bare value (isinstance False, x == Color.RED False), not the member. Correct
the _flatten_enums docstring (was claiming a round-trip that does not happen)
and add test_body_receives_bare_value_not_member, which drives a real wg.run()
and asserts the body sees the bare value - it flips loudly if node-graph later
adds enum reconstruction.

set/frozenset are not descended into: a set fails in general_serializer
regardless of contents (not JSON-serializable, no registered serializer), so
flattening enums inside one would not help. Document this and pin it with
test_flatten_leaves_sets_untouched.

Guard dict-key flattening: two distinct keys collapsing to the same flattened
value (an Enum member and its bare value, or two members sharing a .value) now
raises instead of silently dropping an entry.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: elinscott

src/aiida_workgraph/serialization.py

@@ -12,19 +12,43 @@
 def _flatten_enums(value: Any) -> Any:
     """Replace ``Enum`` members with their bare ``.value``, recursively.
 
-    ``node_graph`` declares that an enum-typed socket's serialized form is
-    its bare value (``socket_spec`` records ``structured_type`` extras so
-    ``coerce_inputs_from_spec`` can rebuild the member before a task body
-    runs). The write path has to honour that contract: without this,
-    ``serialize_ports`` hands the raw ``Enum`` instance to
-    ``general_serializer``, which has no serializer for it and fails at
-    submit time. ``isinstance`` also matches wrapt proxies (TaggedValue)
-    around enum members.
+    ``general_serializer`` (aiida-pythonjob) has no serializer for ``Enum``
+    members: a raw ``Enum`` -- or one nested in a dict/list/tuple -- reaches
+    it and the whole submission fails at submit time with an opaque
+    serialization error. Collapsing members to their ``.value`` here keeps
+    the payload JSON-serializable so ``serialize_ports`` succeeds. The
+    ``isinstance`` check also matches wrapt proxies (node-graph's
+    ``TaggedValue``) wrapping an enum member.
+
+    This is a one-way flattening, *not* a round-trip. Under the pinned
+    ``node_graph`` the read side (``coerce_inputs_from_spec``) rebuilds only
+    structured *models* -- dataclass/pydantic/TypedDict -- and records no
+    ``structured_type`` extra for ``Enum`` sockets, so it never reconstructs
+    the member. Task bodies therefore receive the bare value, not the
+    ``Enum`` they declared. This is harmless for ``IntEnum`` / ``str, Enum``
+    consumed by value or equality, but a body that tests ``x is Color.RED``
+    or ``x == Color.RED`` against a *plain* ``Enum`` sees ``False``. See
+    ``tests/test_serializer.py::test_body_receives_bare_value_not_member``.
+
+    ``set``/``frozenset`` are deliberately left untouched: a set fails in
+    ``general_serializer`` regardless of its contents (not JSON-serializable,
+    no registered serializer), so descending into one to flatten enums would
+    not make it serializable -- see
+    ``tests/test_serializer.py::test_flatten_leaves_sets_untouched``.
     """
     if isinstance(value, Enum):
         return _flatten_enums(value.value)
     if isinstance(value, dict):
-        return {_flatten_enums(k): _flatten_enums(v) for k, v in value.items()}
+        out: Dict[Any, Any] = {}
+        for k, v in value.items():
+            flat_key = _flatten_enums(k)
+            if flat_key in out:
+                # Two distinct keys (e.g. an Enum member and its bare value,
+                # or two members sharing a ``.value``) collapse to one; raise
+                # rather than silently drop an entry.
+                raise ValueError(f'Enum key flattening collision: multiple keys map to {flat_key!r}')
+            out[flat_key] = _flatten_enums(v)
+        return out
     if isinstance(value, list):
         return [_flatten_enums(v) for v in value]
     if isinstance(value, tuple):

tests/test_serializer.py

@@ -94,6 +94,30 @@ def test_flatten_wrapt_proxied_enum_member():
     assert _flatten_enums({proxied: proxied}) == {'red': 'red'}
 
 
+def test_flatten_dict_key_collision_raises():
+    """Two distinct keys collapsing to the same flattened key raises rather than
+    silently dropping an entry. Here ``Color.RED`` (a plain Enum, so distinct in
+    identity and hash from the bare string) and ``'red'`` are separate keys that
+    both flatten to ``'red'``."""
+    payload = {Color.RED: 1, 'red': 2}
+    assert len(payload) == 2  # distinct keys before flattening
+    with pytest.raises(ValueError, match='collision'):
+        _flatten_enums(payload)
+
+
+def test_flatten_leaves_sets_untouched():
+    """``set``/``frozenset`` are not descended into: they fail in
+    ``general_serializer`` regardless of contents, so flattening enums inside
+    them would not make them serializable. Pin that they pass through unchanged
+    (still holding Enum members)."""
+    s = {Color.RED, Color.BLUE}
+    out = _flatten_enums(s)
+    assert out is s
+    assert Color.RED in out
+    fs = frozenset({Priority.LOW})
+    assert _flatten_enums(fs) is fs
+
+
 def test_serialize_ports_accepts_enum_value(aiida_profile):
     """End-to-end through the adapter: an enum-valued namespace entry serializes
     without raising, landing as the flattened bare value (needs a profile so the
@@ -105,3 +129,29 @@ def test_serialize_ports_accepts_enum_value(aiida_profile):
     out = AiidaSerializationAdapter().serialize_ports({'c': Color.RED, 'n': 3}, spec, store=False)
     assert isinstance(out['c'], orm.Str)
     assert out['c'].value == 'red'
+
+
+@task(outputs=['type_name', 'is_enum', 'eq_member'])
+def _observe_enum(c: Color) -> dict:
+    """Report what the body actually receives for an Enum-typed input."""
+    return {
+        'type_name': type(c).__name__,
+        'is_enum': isinstance(c, Color),
+        'eq_member': c == Color.RED,
+    }
+
+
+def test_body_receives_bare_value_not_member(aiida_profile):
+    """Pin the real round-trip behaviour under the pinned ``node_graph``: a task
+    body declaring a *plain* ``Enum`` input receives the BARE value, not the
+    member. ``node_graph`` records ``structured_type`` extras only for
+    dataclass/pydantic/TypedDict, so ``coerce_inputs_from_spec`` never rebuilds
+    an ``Enum``. If node-graph later adds enum reconstruction, this test flips
+    loudly and the docstring/PR claim must be revisited."""
+    wg = WorkGraph('enum_roundtrip')
+    t = wg.add_task(_observe_enum, name='observe', c=Color.RED)
+    wg.run()
+    assert wg.state == 'FINISHED'
+    assert t.outputs['type_name'].value.value == 'str'
+    assert t.outputs['is_enum'].value.value is False
+    assert t.outputs['eq_member'].value.value is False