Flatten Enum members to bare values on the serialize path

[elinscott]

Problem

node-graph's socket spec records structured_type extras for an Enum-typed socket so its own coercion can rebuild the member before a task body runs — the declared contract is that the serialized form is the bare value. This package's AiiDA adapter never honoured that on the way out: a raw Enum instance reached aiida-pythonjob's general_serializer, which has no serializer registered for it, and failed the whole submission.

class Color(str, Enum):
    RED = "red"

@task()
def observe(c: Color): ...

wg = WorkGraph()
wg.add_task(observe, c=Color.RED)
wg.run()
# before this fix: submission fails — general_serializer has no case for Color

Changes

• _flatten_enums walks a value recursively — dict keys and values, list/tuple items, nested structures — and replaces every Enum member with its bare value before serialize_ports hands the payload to aiida-pythonjob. set/frozenset are left alone and documented as out of scope: a set fails general_serializer regardless of its contents (no registered serializer, not JSON-safe), so flattening inside one wouldn't help.
• An enum-free payload comes back as the exact object passed in, not a rebuilt copy: rebuilding unconditionally had been downgrading a namedtuple to a plain tuple and an OrderedDict/defaultdict to a plain dict even when nothing needed flattening. A container that does carry an Enum still flattens as before.
• Dict-key flattening is guarded: two distinct keys that collapse to the same flattened value (an Enum member and its own bare value, or two members sharing a .value) now raises instead of silently dropping an entry.
• The read side — whether a task body gets the bare value back or the reconstructed member — is node-graph's call, not this package's (paired PR 🐛 Decide enum/Literal socket membership once, at assignment scinode/node-graph#178, which closes Literal[...] not supported scinode/node-graph#175 and answers Change of syntax WorkGraph.add_link #176 for annotated sockets). With node-graph 0.6.5 an Enum-typed input arrives as the flattened value; with Change the syntax for wait #178 it arrives as the member. A body relies on its annotation either way and never rebuilds a member by hand from what arrived. The tests here pin that agreement: what arrives matches what the installed node-graph advertises, so a reconstruction that silently regresses fails here rather than passing vacuously.

Testing

• tests/test_serializer.py: _flatten_enums unit coverage (bare/IntEnum/str-Enum members, dict keys and values, list/tuple, mixed nesting, enum-free passthrough, namedtuple/OrderedDict/defaultdict type preservation on the enum-free path and continued flattening when they do carry an Enum, wrapt-proxied members) plus an end-to-end serialize_ports check that an enum-valued entry serializes; a live wg.run() exercising a function task with an Enum input, asserting the arriving form against the installed node-graph's capability, plus a @task.graph body using the member as annotated (skipped where node-graph does not reconstruct).
• tests/test_serializer.py against node-graph 0.6.5 (the pinned release): 16 passed, 1 skipped — the @task.graph member test skips because 0.6.5 does not reconstruct Enum sockets, and the arrival test asserts the flattened form.
• The same file under a PYTHONPATH shadow of the paired node-graph branch (verified via node_graph.__file__): 17 passed, 0 skipped — the member arrives and the graph body uses c.name as annotated. The pair of runs is the discriminating check: the same tests assert opposite arrival forms and pass under each node-graph exactly as its capability predicts.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f59c251f-21d6-4e76-a2c8-31b9ada7a4ff

📥 Commits

Reviewing files that changed from the base of the PR and between ce3ca42 and 885fdb1.

📒 Files selected for processing (2)

• src/aiida_workgraph/serialization.py
• tests/test_serializer.py

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.

---

📝 Walkthrough

Walkthrough

Enum values now flatten recursively in dictionaries, lists, and tuples before port serialization. Unchanged containers retain their instances. Dictionary key collisions still raise ValueError. Tests cover container preservation and capability-dependent WorkGraph Enum delivery.

Changes

Enum serialization

Layer / File(s)	Summary
Recursive Enum flattening src/aiida_workgraph/serialization.py	_flatten_enums flattens Enum keys and values in supported containers. serialize and serialize_ports apply the conversion before port serialization.
Container preservation validation tests/test_serializer.py	Tests verify that enum-free namedtuples and dictionary subclasses remain unchanged, while Enum-containing containers are flattened.
WorkGraph Enum reconstruction tests/test_serializer.py	Tests distinguish Enum-member delivery from flattened-string delivery and gate graph-body checks on node_graph capability.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 885fd

The change recursively converts enum values during serialization and adds coverage for supported input forms; no actionable merge-blocking risk remains after normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant WorkGraphTask
  participant serialization.serialize
  participant _flatten_enums
  participant serialization.serialize_ports
  WorkGraphTask->>serialization.serialize: provide Enum input
  serialization.serialize->>_flatten_enums: flatten nested Enum values
  _flatten_enums-->>serialization.serialize: return converted or original container
  serialization.serialize->>serialization.serialize_ports: pass flattened payload
  serialization.serialize_ports-->>WorkGraphTask: provide serialized port data

Loading

Suggested reviewers: geigerj2

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: flattening Enum members to their bare values during serialization.
Description check	✅ Passed	The description directly explains the Enum serialization problem, implementation, limitations, collision behavior, and test coverage.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 25 functions across 2 files.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.02%. Comparing base (46148af) to head (ce3ca42).
⚠️ Report is 24 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #800      +/-   ##
==========================================
+ Coverage   90.18%   91.02%   +0.85%     
==========================================
  Files          44       46       +2     
  Lines        2991     3184     +193     
==========================================
+ Hits         2697     2898     +201     
+ Misses        294      286       -8     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.