Skip to content

Flatten Enum members to bare values on the serialize path#800

Open
elinscott wants to merge 3 commits into
aiidateam:mainfrom
elinscott:serialize-enum-flatten
Open

Flatten Enum members to bare values on the serialize path#800
elinscott wants to merge 3 commits into
aiidateam:mainfrom
elinscott:serialize-enum-flatten

Conversation

@elinscott

@elinscott elinscott commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

Problem

aiida-pythonjob's general_serializer has no serializer for Enum members. A raw Enum instance (or a dict/list containing one, e.g. a parameters dict with an enum-valued entry) flows straight into serialize_ports, which hands it to general_serializer, and the whole submission fails at submit time with an opaque serialization error.

Change

A small _flatten_enums helper replaces Enum members with their .value, recursively through dicts (keys and values), lists and tuples, and is applied in both AiidaSerializationAdapter.serialize and AiidaSerializationAdapter.serialize_ports immediately after resolve_tagged_values. The isinstance(value, Enum) check also matches wrapt proxies (TaggedValue) wrapping enum members, so tagged graph inputs are handled by the same path.

This is a one-way flattening on the write path, not a round-trip. Under the pinned node-graph (0.6.5) the read side (coerce_inputs_from_spec) reconstructs only structured models — dataclass/pydantic/TypedDict — and records no structured_type extra for Enum sockets, so it never rebuilds 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 / x == Color.RED against a plain Enum will see False. True reconstruction would need an enum branch in node-graph's structured_type_info/coerce_structured_value and is left as follow-up there.

Two edge cases are handled explicitly: set/frozenset are 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; documented as a limitation), and dict-key flattening now raises on a collision (two distinct keys collapsing to one, e.g. an Enum member and its bare value) rather than silently dropping an entry.

Testing

New unit tests in tests/test_serializer.py, including test_body_receives_bare_value_not_member, which drives a real wg.run() and asserts a plain-Enum input reaches the body as the bare value — pinning the actual (non-round-trip) behaviour so it flips loudly if node-graph later adds enum reconstruction — plus tests for the set passthrough and the key-collision guard.

Exercised end-to-end by the downstream Koopmans package, whose workgraphs pass enum-valued settings (e.g. functional and spin-treatment enums) through @task.graph inputs: before the change every such submission failed in general_serializer; after it, the graphs submit (the bodies receive the bare enum value, per the note above).

elinscott and others added 2 commits July 8, 2026 13:46
node-graph's socket spec records structured_type extras for enum-typed
sockets so coerce_inputs_from_spec can rebuild the member before a task
body runs - the declared contract is that the serialized form is the
bare value. The AiiDA adapter never honoured it: raw Enum instances
reached aiida-pythonjob's general_serializer, which has no serializer
for them and failed the whole submission. Flatten enums (recursively,
including dict keys/values and list items) before serialize_ports.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
node-graph's socket spec records structured_type extras for enum-typed
sockets so coerce_inputs_from_spec can rebuild the member before a task
body runs - the declared contract is that the serialized form is the
bare value. The AiiDA adapter never honoured it: raw Enum instances
reached aiida-pythonjob's general_serializer, which has no serializer
for them and failed the whole submission. Flatten enums (recursively,
including dict keys/values and list items) before serialize_ports.

Add unit tests in tests/test_serializer.py covering _flatten_enums
(bare/IntEnum/str-Enum members, dict keys and values, list/tuple,
mixed nesting, enum-free passthrough, wrapt-proxied members) plus one
end-to-end serialize_ports check that an enum-valued entry serializes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@elinscott
elinscott marked this pull request as ready for review July 8, 2026 12:46
…ation

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>
@elinscott
elinscott force-pushed the serialize-enum-flatten branch from 3a86513 to e53a7ce Compare July 17, 2026 09:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant