Allow UserStepDecorators to also be StepMutators#3079
Conversation
Greptile SummaryThis PR enables a class to simultaneously inherit from both Confidence Score: 4/5Safe to merge; all previously-flagged P0/P1 issues are resolved in this revision The remaining complexity is in the dual-inheritance lifecycle (five call sites for external_init across three files) which is correctly guarded by _ran_init, but the non-trivial interaction surface warrants a 4 rather than 5 metaflow/user_decorators/user_step_decorator.py (add_or_raise / _all_step_fields) and metaflow/user_decorators/mutable_step.py (add_decorator string-path) Important Files Changed
Reviews (17): Last reviewed commit: "add_decorator IGNORE path: return None i..." | Re-trigger Greptile |
12f3a80 to
86e5908
Compare
518e3c0 to
b698ae5
Compare
| def init(self, marker="default"): | ||
| self._marker = marker | ||
|
|
||
| def mutate(self, mutable_step): |
There was a problem hiding this comment.
could also cover pre_mutate, but validating one or the other should be enough of a confirmation as there's no cross-talk between the two?
There was a problem hiding this comment.
Yes, I could add both but yes, it should not have cross talk.
| return False | ||
|
|
||
|
|
||
| def _init(flow, only_non_static=False): |
There was a problem hiding this comment.
so if an extension has been relying on attaching decorators, and needing to init these, instead of decorators._init(), will deco.external_init() now serve the same purpose?
There was a problem hiding this comment.
So yes, this function just called all the external_init functions. I now spread out the external_init in different parts so that it is more accurate about when to call them.
| # These decospecs are the ones from run/resume/spin PLUS the ones from the | ||
| # environment (for example the @conda) | ||
| decorators._attach_decorators(obj.flow, all_decospecs) | ||
| decorators._init(obj.flow) |
There was a problem hiding this comment.
why is an init not required anymore?
There was a problem hiding this comment.
It's still there but the calls are spread out.
| "Evaluating step level decorator %s for %s (pre-mutate)" | ||
| % (deco.__class__.__name__, step.name) | ||
| ) | ||
| deco.pre_mutate( |
There was a problem hiding this comment.
just a review note: this is fine, because we check in the previous loop that all decos in the config_decorators are StepMutators.
What is a bit unclear is why there was a need to split these to two loops. Is calling external_init() on every config deco before mutating a requirement?
There was a problem hiding this comment.
I was trying to keep with the other hooks where all hooks are called first and then all the second one.
Also minor changes in error messages
- Clean up "hooks" for mutators - Allow flow mutators to add other flow mutators - Fix an issue with step mutators adding step mutators - Return the decorator added with add_decorator
Allows looking up a step's mutable proxy by name without iterating through all steps. Raises MetaflowException if the step doesn't exist. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- has_decorator(name): check if a step has a decorator by name - get_decorator_kwargs(name): get the kwargs of the first matching decorator Enables cleaner decorator inspection without manual tuple iteration. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Merge-and-replace a decorator's kwargs in one call, with optional reject_if for declaring incompatible options. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Returns all decorator kwargs for a given name, enabling grouped queries without manual iteration. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Called automatically during step_init to validate this decorator against all others on the same step. Subclasses override to declare incompatible combinations. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Runtime equivalent of MutableStep.get_decorator_kwargs — finds a decorator's attributes on an instantiated step, with optional exclude filter for skipping specific attribute values. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Same query API as MutableStep, but for flow-level decorators. Eliminates manual iteration of flow.decorator_specs. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…s more consistent
ca4d586 to
52d5587
Compare
The error message referenced deco_type.name which doesn't exist on FlowMutator classes (they use decorator_name). This caused an opaque AttributeError instead of the intended MetaflowException when calling add_decorator(SomeFlowMutator) from mutate(). Use getattr fallback chain: .name -> .decorator_name -> repr.
The dual_deco test now exercises all four lifecycle methods for a class inheriting from both UserStepDecorator and StepMutator: pre_mutate, mutate, pre_step, and post_step.
get_decorator_by_name returns None for unregistered names, which caused issubclass(None, ...) to raise TypeError. Add a None guard so unregistered names fall through to the StepDecorator name-based removal path gracefully.
- CRITICAL-1: Call external_init() before appending to lists so a failure doesn't leave a half-initialized mutator in the decorator registries. - MAJOR-1: OVERRIDE mode now uses slice-assignment (merged[:] = ...) instead of list.remove() to avoid skipping elements in the outer pre_mutate iteration loop that holds a reference to the same list.
The name-based caching of skipped decorators would cause all instances of an allow_multiple decorator (e.g. @card) to be skipped if any single instance was skipped for spin mode. Revert to per-decorator checking, matching the behavior on master.
The flow_fixture factory calls request.config.getoption("--use-latest")
but the option was never registered via pytest_addoption in this
conftest. It worked when run from a parent conftest that registered
the option, but broke when the spin tests were invoked from a batch
runner that doesn't pull in that parent.
Register the option locally and pass default=False to getoption() for
belt-and-suspenders safety.
saikonen
left a comment
There was a problem hiding this comment.
extension changes verified to work, should be good to go from that side.
Previously `remove_decorator(name)` called without deco_args/deco_kwargs (do_all=True) would skip the `_pre_mutate + isinstance(deco, StepMutator)` check via its early `continue`, letting a user strip a StepMutator from `mutate()` in violation of the documented contract. Restructure the match/guard so both the do_all path and the specific-match path go through the same StepMutator-from-mutate check. Add a unit test covering the guard in both directions.
When a UserStepDecorator's post_step returned (None, False), task.py
would unpack the tuple (raised_exception=None, fake_next_call_args=False),
then reach the "Invalid value passed to self.next" RuntimeError branch
because False matches neither True, {}, nor had_raised_exception.
False is the natural "don't override self.next" return shape, semantically
identical to None. Accept both at the early guard so (None, False) is a
no-op as users reasonably expect.
Add a regression flow + tests exercising this return shape.
When `MutableStep.add_decorator` handles a UserStepDecoratorBase and `add_or_raise` suppresses the registration (duplicates=IGNORE, or non-statically-defined ERROR), the decorator was never actually added to the step. But the outer `return step_deco` / `return d` was still handing back the throwaway instance, contradicting the documented "None if not added due to duplicate handling" contract and diverging from the StepDecorator path which correctly returns None via the added_deco tracker. Make `add_or_raise` return `self` when the decorator was registered and `None` when the call was swallowed. Propagate that signal at both UserStepDecoratorBase call sites (string-form and class-form) so the public return value matches the contract and the StepDecorator path. Also fix _all_step_fields non-determinism: it used a set for dedup and returned `list(set)`, giving non-deterministic ordering. Switched to a parallel ordered list + dedup set so callers see a deterministic MRO order (matters for side-effecting traversals like OVERRIDE cleanup).
## Summary
Allow a class to inherit from both `UserStepDecorator` and `StepMutator`
at the same time, so a single decorator can contribute both a
`mutate()`/`pre_mutate()` pass over the flow/step (the `StepMutator`
role) and `pre_step`/`post_step` wrapping hooks (the `UserStepDecorator`
role). The instance is registered on both `step.wrappers` and
`step.config_decorators` via a new MRO-walking `_all_step_fields()`
helper, with an `_ran_init` idempotency guard so `external_init()` is
not called twice.
In the process we also tightened a number of nearby corners in the
user-decorator / flow-decorator plumbing and added several public
helpers that made the dual-inheritance pattern writable at all.
## PR Type
- [x] Bug fix
- [x] New feature
- [ ] Core Runtime change
- [ ] Docs / tooling
- [ ] Refactoring
## New functionality
**Dual inheritance** — `class foo(UserStepDecorator, StepMutator)` is
now a first-class pattern:
- Registration is unified via `_all_step_fields()` which walks the MRO
and returns every `_step_field` declared on ancestor classes
(deterministic MRO order).
- `_ran_init` guard on `external_init()` ensures `init()` is only
invoked once per instance regardless of how many fields the instance
lives on.
- `_init_step_decorators` now calls `external_init` on `step.wrappers`
too (previously only `step.decorators`); `_process_config_decorators`
covers `config_decorators`. This fixes a pre-existing gap where a pure
`UserStepDecorator` wrapper's `init()` could be skipped.
**New MutableStep / MutableFlow helpers** (to support writing
dual-inherit decorators without reaching into internal lists):
- `MutableFlow.get_step(name)`
- `MutableFlow.has_decorator(name)` / `MutableStep.has_decorator(name)`
- `MutableFlow.get_decorator_kwargs(...)` /
`MutableStep.get_decorator_kwargs(...)`
- `MutableStep.get_all_decorator_kwargs(...)`
- `MutableStep.replace_decorator(...)`
- `MutableStep.decorator_specs` dedupes instances that live on multiple
step fields.
**New StepDecorator hooks:**
- `check_step_conflicts(decorators)` — called during
`_init_step_decorators` so a decorator can reject incompatible neighbors
early.
- `find_decorator_attrs(...)` convenience for decorator-to-decorator
introspection.
**Task-runtime plumbing:**
- `task.py` post-step sentinel changed from `False` to `None`, with
explicit `is not False` guard so `post_step` returning `(None, False)`
is a no-op instead of tripping `"Invalid value passed to self.next"`.
- `cli.py` no longer keeps the old manual `_init()` / `external_init()`
loops for flow mutators — it's all driven by
`_process_config_decorators` and `_init_step_decorators` now.
- The standalone `decorators._init()` helper was removed; its
responsibilities were absorbed into `_init_flow_decorators` and
`_init_step_decorators`.
## Bug fixes (most flagged by Greptile during review)
- **`add_decorator` on FlowMutator** — fixed `AttributeError:
deco_type.name` when `add_decorator(MyFlowMutator)` is called from
`mutate()`; the error path now uses `getattr(deco_type, "name", None) or
getattr(deco_type, "decorator_name", deco_type)` so the real
`MetaflowException` is raised.
- **`remove_decorator("typo")`** — `get_decorator_by_name` returns
`None` for unregistered names; no longer passes that through to
`issubclass(None, ...)` and raises `TypeError`.
- **FlowMutator registration** — `_do_add` now calls `external_init()`
**before** appending to the decorator lists, so a raising
`external_init` doesn't leave a half-initialized mutator in `_self_data`
/ merged mutators.
- **OVERRIDE + merged list** — switched from `list.remove()` (which
shifts indices under the outer `pre_mutate` iteration) to
slice-assignment `merged[:] = [...]`. No more silently-skipped elements.
- **`cached_skip_decorators`** — removed a name-based cache that
over-skipped `allow_multiple` decorators (e.g. multiple `@card`s on
different steps). Matches master's per-decorator behavior.
- **`remove_decorator(name)` `do_all=True`** — the `_pre_mutate +
isinstance(StepMutator)` guard is now enforced on both the
specific-match and `do_all` paths; previously `do_all` could strip a
`StepMutator` from `mutate()` in violation of the documented contract.
- **`add_decorator` IGNORE path** — the `UserStepDecoratorBase` branches
now return `None` when `add_or_raise` swallowed the registration
(duplicate + IGNORE, or non-statically-defined ERROR), matching the
`StepDecorator` path and the documented contract. Previously the
throwaway instance was returned as if it had been added.
## Tests
- New `test/unit/mutators/` suite with dedicated flows + tests for dual
inheritance, dynamic flow-mutator addition, string-form mutator
addition, `add_decorator` return value, and regression coverage for the
fixes above (`test_remove_decorator_guard.py`,
`test_post_step_none_false.py`).
- Dual inheritance test exercises all four lifecycle methods:
`pre_mutate`, `mutate`, `pre_step`, `post_step`.
- Spin conftest registers `--use-latest` so the fixture factory works
when the suite is driven by a batch runner that doesn't inherit the
option from a parent conftest.
---------
Co-authored-by: Nissan Pow <npow@netflix.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Summary
Allow a class to inherit from both
UserStepDecoratorandStepMutatorat the same time, so a single decorator can contribute both amutate()/pre_mutate()pass over the flow/step (theStepMutatorrole) andpre_step/post_stepwrapping hooks (theUserStepDecoratorrole). The instance is registered on bothstep.wrappersandstep.config_decoratorsvia a new MRO-walking_all_step_fields()helper, with an_ran_initidempotency guard soexternal_init()is not called twice.In the process we also tightened a number of nearby corners in the user-decorator / flow-decorator plumbing and added several public helpers that made the dual-inheritance pattern writable at all.
PR Type
New functionality
Dual inheritance —
class foo(UserStepDecorator, StepMutator)is now a first-class pattern:_all_step_fields()which walks the MRO and returns every_step_fielddeclared on ancestor classes (deterministic MRO order)._ran_initguard onexternal_init()ensuresinit()is only invoked once per instance regardless of how many fields the instance lives on._init_step_decoratorsnow callsexternal_initonstep.wrapperstoo (previously onlystep.decorators);_process_config_decoratorscoversconfig_decorators. This fixes a pre-existing gap where a pureUserStepDecoratorwrapper'sinit()could be skipped.New MutableStep / MutableFlow helpers (to support writing dual-inherit decorators without reaching into internal lists):
MutableFlow.get_step(name)MutableFlow.has_decorator(name)/MutableStep.has_decorator(name)MutableFlow.get_decorator_kwargs(...)/MutableStep.get_decorator_kwargs(...)MutableStep.get_all_decorator_kwargs(...)MutableStep.replace_decorator(...)MutableStep.decorator_specsdedupes instances that live on multiple step fields.New StepDecorator hooks:
check_step_conflicts(decorators)— called during_init_step_decoratorsso a decorator can reject incompatible neighbors early.find_decorator_attrs(...)convenience for decorator-to-decorator introspection.Task-runtime plumbing:
task.pypost-step sentinel changed fromFalsetoNone, with explicitis not Falseguard sopost_stepreturning(None, False)is a no-op instead of tripping"Invalid value passed to self.next".cli.pyno longer keeps the old manual_init()/external_init()loops for flow mutators — it's all driven by_process_config_decoratorsand_init_step_decoratorsnow.decorators._init()helper was removed; its responsibilities were absorbed into_init_flow_decoratorsand_init_step_decorators.Bug fixes (most flagged by Greptile during review)
add_decoratoron FlowMutator — fixedAttributeError: deco_type.namewhenadd_decorator(MyFlowMutator)is called frommutate(); the error path now usesgetattr(deco_type, "name", None) or getattr(deco_type, "decorator_name", deco_type)so the realMetaflowExceptionis raised.remove_decorator("typo")—get_decorator_by_namereturnsNonefor unregistered names; no longer passes that through toissubclass(None, ...)and raisesTypeError._do_addnow callsexternal_init()before appending to the decorator lists, so a raisingexternal_initdoesn't leave a half-initialized mutator in_self_data/ merged mutators.list.remove()(which shifts indices under the outerpre_mutateiteration) to slice-assignmentmerged[:] = [...]. No more silently-skipped elements.cached_skip_decorators— removed a name-based cache that over-skippedallow_multipledecorators (e.g. multiple@cards on different steps). Matches master's per-decorator behavior.remove_decorator(name)do_all=True— the_pre_mutate + isinstance(StepMutator)guard is now enforced on both the specific-match anddo_allpaths; previouslydo_allcould strip aStepMutatorfrommutate()in violation of the documented contract.add_decoratorIGNORE path — theUserStepDecoratorBasebranches now returnNonewhenadd_or_raiseswallowed the registration (duplicate + IGNORE, or non-statically-defined ERROR), matching theStepDecoratorpath and the documented contract. Previously the throwaway instance was returned as if it had been added.Tests
test/unit/mutators/suite with dedicated flows + tests for dual inheritance, dynamic flow-mutator addition, string-form mutator addition,add_decoratorreturn value, and regression coverage for the fixes above (test_remove_decorator_guard.py,test_post_step_none_false.py).pre_mutate,mutate,pre_step,post_step.--use-latestso the fixture factory works when the suite is driven by a batch runner that doesn't inherit the option from a parent conftest.