Skip to content

Refactor Python unit context handling into ContextItem and UnitContext#148

Draft
Copilot wants to merge 2 commits into
chore/SOF-7921from
copilot/refactor-context-item-handling
Draft

Refactor Python unit context handling into ContextItem and UnitContext#148
Copilot wants to merge 2 commits into
chore/SOF-7921from
copilot/refactor-context-item-handling

Conversation

Copilot AI commented Jun 3, 2026

Copy link
Copy Markdown
Contributor

PR #147 had context-item parsing/serialization and list mutation logic embedded in ExecutionUnit, which made the unit class carry provider-shape normalization concerns. This refactor aligns the Python implementation with the TypeScript structure by separating single-item and collection responsibilities.

  • Context item model extracted

    • Added src/py/mat3ra/wode/units/context_item.py with ContextItem as the canonical persisted shape:
      • {"name", "isEdited", "data", "extraData"}
    • Centralized normalization for:
      • persisted dict input (name/isEdited/data/extraData)
      • provider-yield input (single provider key + optional isXxxEdited / xxxExtraData)
      • ignored isUsingJinjaVariables
    • Moved scalar-data wrapping (data -> {"value": data}) and unwrapping ({"value": ...} -> ...) into the class.
  • Unit context collection extracted

    • Added src/py/mat3ra/wode/units/unit_context.py with UnitContext (RootModel[List[ContextItem]]).
    • UnitContext now owns list-level operations:
      • get, get_data, add/upsert, remove, clear, iteration/length.
    • Replace-by-name semantics are implemented in one place.
  • ExecutionUnit slimmed down

    • context field switched to UnitContext.
    • Field coercion delegates to UnitContext.from_value(...).
    • Removed in-class context helpers/parsers:
      • _context_item_name, _read_context_data, context_item, _replace_context_item, _normalized_context_item, _context_item_from_provider_yield.
    • Context methods are now thin pass-throughs:
      • add_context, get_context_item, get_context, set_context, remove_context, clear_context.
    • add_context_provider(provider) behavior remains unchanged.
  • Exports and compatibility

    • Re-exported ContextItem and UnitContext from src/py/mat3ra/wode/units/__init__.py.
    • Preserved existing call-site behavior in tests and workflow/subworkflow code:
      • unit.add_context(...) from persisted or provider-yield input
      • unit.get_context(...) data semantics (including scalar unwrap)
      • serialized to_dict()["context"] remains webapp-compatible list-of-dicts shape.
unit.add_context({"degauss": 0.001, "isDegaussEdited": False, "degaussExtraData": {"units": "Ry"}})

unit.get_context("degauss")
# -> 0.001

unit.get_context_item("degauss")
# -> {"name": "degauss", "isEdited": False, "data": {"value": 0.001}, "extraData": {"units": "Ry"}}
Original prompt

Context

PR #147 currently keeps the entire context-item handling logic inside ExecutionUnit (see src/py/mat3ra/wode/units/execution.py), with many static helpers (_context_item_name, _read_context_data, context_item, _replace_context_item, _normalized_context_item, _context_item_from_provider_yield, plus add_context/get_context/set_context/remove_context/clear_context).

Reviewer feedback on #147 (comment):

The implementation should be simpler. Dealing with context item serialization should be in a class of its own; the list of context items should have its own logic.

This mirrors how the TypeScript implementation is structured. We should align the Python side with the TS side.

TypeScript reference implementation

Relevant TS files at https://github.com/Exabyte-io/wode/tree/main/src/js :

src/js/context/providers/base/ContextProvider.ts:

  • Defines UnitContext = ContextItemSchema[], ContextName, ContextExtraData, ContextData.
  • ContextProvider.getContextItemData(): S returns { name, isEdited, data, extraData } — the canonical serialization of a single persisted context item. The provider owns serialization of itself into a persisted item.
  • ContextProvider.findContextItem(unitContext, contextName) is a static helper that finds an item in a UnitContext array by name.

src/js/units/ExecutionUnit.ts:

  • context is just ContextItemSchema[] (a plain list).
  • ExecutionUnit has NO add_context/get_context/remove_context/context_item helpers, NO yield-data parsing. It simply does:
    savePersistentContext() {
        const persistentItems = this.contextProvidersInstances.map((p) => p.getContextItemData());
        this.context = persistentItems.filter((c) => c.isEdited);
    }
  • All "is this edited / is this a kgrid / what's the extraData" logic lives on the provider class, not on the unit.

Goal

Refactor the Python implementation so that:

  1. A ContextItem class of its own handles serialization of a single context item. It should wrap/produce the canonical persisted shape:

    {"name": str, "isEdited": bool, "data": Any, "extraData": Dict[str, Any]}
    • It should be constructable from:
      • the persisted dict shape ({"name", "isEdited", "data", "extraData"}), and
      • a "yielded" provider dict (the _context_item_from_provider_yield logic — single provider key + optional isXxxEdited / xxxExtraData siblings, ignoring isUsingJinjaVariables).
    • It should validate/normalize once, in this class — not inside ExecutionUnit.
    • It should be JSON-serializable via Pydantic (so it round-trips inside ExecutionUnitSchema's context field). Prefer a Pydantic model (e.g. subclass ContextItemSchema from ESSE, mirroring how TS uses ContextItemSchema).
    • The data value passed in as a non-dict should be wrapped as {"value": data} exactly as the current context_item(...) helper does, and _read_context_data (which unwraps {"value": ...}) should live on this class too.
  2. A UnitContext collection class (or thin wrapper) owns list-level logic. It should expose only the operations the unit actually needs, e.g.:

    • get(name) -> Optional[ContextItem]
    • get_data(name, default=None) -> Any (equivalent of current ExecutionUnit.get_context)
    • add(item) / upsert(item) (replace-by-name semantics from _replace_context_item)
    • remove(name)
    • clear()
    • iteration / __len__ / __iter__
    • JSON serialization back to List[ContextItemSchema] so the Pydantic field round-trips unchanged.

    Pick whichever shape is most idiomatic in Pydantic v2 — either a plain class UnitContext(RootModel[List[ContextItem]]) or a small standalone class that's converted via a field_validator/field_serializer. The goal is that ExecutionUnit.context is that type and all list manipulation happens on it.

  3. Slim ExecutionUnit back down. After the refactor, ExecutionUnit should ideally contain only:

    • the context: UnitContext field (with a field_validator that delegates to UnitContext/ContextItem for coercion from raw dicts / lists, replacing the current _validate_context),
    • thin pass-through methods needed by callers: add_context(item), get_context_item(name), get_context(name, default), set_context(items), remove_context(name), clear_context() — each one line that delegates to self.context.
    • add_context_provider(provider) continues to work by calling self.add_context(provider.yield_data()).

    Remove from ExecutionUnit: _context_item_name, _read_context_data, context_item, _replace_context_item, _normalized_context_item, _context_item_from_provider_yield. They belong on ContextItem / UnitContext.

  4. Module placement. Put the new classes in a dedicated module, e.g.:

    • src/py/mat3ra/wode/units/context_item.pyContextItem
    • `src/py/mat3ra/wo...

This pull request was created from Copilot chat.

Copilot AI changed the title [WIP] Refactor context-item handling logic into its own class Refactor Python unit context handling into ContextItem and UnitContext Jun 3, 2026
Copilot AI requested a review from timurbazhirov June 3, 2026 04:17
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.

2 participants