feat: context-manager form of Container.override via OverrideHandle (INT-4)

architecture/testing-and-overrides.md

@@ -14,13 +14,19 @@ tree — all child containers hold a reference to the same registry instance.
 ### container.override and container.reset_override
 
 ```python
-container.override(provider: AbstractProvider[T], override_object: T) -> None
+container.override(provider: AbstractProvider[T], override_object: T) -> OverrideHandle[T]
 container.reset_override(provider: AbstractProvider[T] | None = None) -> None
 ```
 
-`container.override(provider, obj)` writes `obj` into the shared `OverridesRegistry` under the provider's id.
-`container.reset_override(provider)` removes that entry. Calling `reset_override()` with no argument (or `None`)
-clears **all** overrides from the registry.
+`container.override(provider, obj)` writes `obj` into the shared `OverridesRegistry` under the provider's id and
+returns an `OverrideHandle[T]`, generic over the override object's type. The override is active from the `override()`
+call itself, not from `__enter__` — imperative callers that discard the handle see identical behavior to before.
+Used as a context manager, the handle's `__exit__` restores the snapshot taken at the `override()` call — the
+provider's prior override if one existed, otherwise no override — unconditionally, even on exception and even if
+`reset_override()` ran inside the block. Nested overrides of the same provider unwind in order, each handle
+restoring what was active before it. The `OverridesRegistry` itself stays a flat dict; the stack lives in the
+handles, not the registry. `container.reset_override(provider)` removes that entry directly. Calling
+`reset_override()` with no argument (or `None`) clears **all** overrides from the registry.
 
 Because the registry is shared, calling either method on a child container has the same effect as calling it on the
 root — the override is visible tree-wide. `close_async` and `close_sync` on the root container also call

docs/introduction/for-fastapi-users.md

@@ -16,6 +16,7 @@ translates the `Depends` idioms you already know into their modern-di equivalent
 | `yield`-based teardown (`def fn(): ...; yield x; ...cleanup...`) | `cache=CacheSettings(finalizer=cleanup_fn)` | modern-di has no generator-creator form (see [Design decisions](design-decisions.md)); teardown is a second, explicit object instead of code after `yield`. `finalizer` may be sync or async — see [Lifecycle](../providers/lifecycle.md). |
 | `@lru_cache`-wrapped dependency (process-wide singleton) | `Factory(fn, scope=Scope.APP, cache=True)`, optionally with a `finalizer` | `lru_cache` has no cleanup hook; the APP-scoped cached `Factory` adds one via `CacheSettings(finalizer=...)` if the singleton needs to release anything on shutdown. |
 | `app.dependency_overrides[fn] = fake` | `container.override(provider, fake)` | modern-di overrides are keyed by **provider reference**, not by callable, and apply across the whole container tree — see [Testing with overrides](../recipes/testing-overrides.md). Reset with `container.reset_override(provider)`. |
+| the manual `try`/`finally` reset FastAPI's docs recommend around `dependency_overrides` | `with container.override(provider, fake) as mock: ...` | Auto-resets on exit instead of a hand-written `finally`. See [Testing with overrides](../recipes/testing-overrides.md) for the full semantics. |
 
 ## Two meanings of "scope"
 

docs/migration/from-dependency-injector.md

@@ -84,7 +84,7 @@ Use this table as the index for the rest of the guide. Every provider class docu
 | `DeclarativeContainer` | `Group` (schema) + `Container(groups=[...], validate=True)` (runtime) | [§2](#2-key-conceptual-shifts) |
 | `container.init_resources()` | Lazy initialization — no equivalent needed | [§9](#9-testing-and-overrides) |
 | `container.shutdown_resources()` / `provider.shutdown()` | `container.close_sync()` / `await container.close_async()` | [§9](#9-testing-and-overrides) |
-| `provider.override(...)` / `with provider.override(...):` | `container.override(provider, mock)` (no context-manager form yet — see [§9](#9-testing-and-overrides)) | [§9](#9-testing-and-overrides) |
+| `provider.override(...)` / `with provider.override(...):` | `container.override(provider, mock)` / `with container.override(provider, mock):` — see [§9](#9-testing-and-overrides) | [§9](#9-testing-and-overrides) |
 | `provider.reset_override()` / `provider.reset_last_overriding()` | `container.reset_override(provider)` | [§9](#9-testing-and-overrides) |
 
 ## 4. Migrate the dependency graph
@@ -306,7 +306,15 @@ container.override(AppGroup.api_client_factory, unittest.mock.Mock(ApiClient))
 container.reset_override(AppGroup.api_client_factory)  # or reset_override() to clear all
 ```
 
-`dependency-injector` also has a context-manager override form (`with container.api_client_factory.override(mock):`) that auto-resets on exit. **`modern-di` does not have this form yet** — `container.override()`/`reset_override()` is imperative-only today; pair the calls manually (e.g. in a pytest fixture's teardown, or `try`/`finally`). See [Testing with overrides](../recipes/testing-overrides.md) for tree-wide sharing and reset mechanics.
+`dependency-injector` also has a context-manager override form (`with container.api_client_factory.override(mock):`) that auto-resets on exit. `modern-di` has the same shape — `with container.override(provider, mock) as m:` applies the override for the block and restores the prior state on exit, including on exception:
+
+```python
+# modern-di
+with container.override(AppGroup.api_client_factory, unittest.mock.Mock(ApiClient)) as mock_factory:
+    ...
+```
+
+See [Testing with overrides](../recipes/testing-overrides.md) for tree-wide sharing, nesting, and reset mechanics.
 
 ### Lifecycle
 

docs/recipes/testing-overrides.md

@@ -4,7 +4,16 @@
 
 ## Solution
 
-`container.override(provider, replacement)` replaces what the provider resolves to. The replacement is keyed by **provider reference** (not name) and is shared across the container tree, so an override on the root APP container applies to all child REQUEST containers too. Reset with `container.reset_override(provider)` (or `container.reset_override()` to clear all).
+`container.override(provider, replacement)` replaces what the provider resolves to, immediately, and returns an `OverrideHandle`. Used as a context manager, it auto-resets on exit — this is the primary spelling for tests:
+
+```python
+with container.override(MyGroup.api_client, mock_client) as client:
+    ...  # resolution returns mock_client; prior state restored on exit
+```
+
+The override applies at the `override()` call, not at `__enter__`. `__exit__` restores the snapshot taken at that call — a previously stacked override if there was one, otherwise no override — even on exception, and even if `reset_override()` — or a root `close_sync()`/`close_async()`, which clears all overrides — ran inside the block; exit still restores the snapshot. Nested overrides of the same provider unwind in order: each handle restores whatever was active before it. Handles are expected to exit in reverse order of creation — `with`-block nesting does this naturally; manually exiting handles out of order can restore stale state.
+
+`container.override(provider, replacement)` also works as a plain imperative call: reset with `container.reset_override(provider)` (or `container.reset_override()` to clear all). This pair remains fully supported — see the patterns below — and `close_sync`/`close_async` on the root container also clear all overrides automatically. Either way, the replacement is keyed by **provider reference** (not name) and is shared across the container tree, so an override on the root APP container applies to all child REQUEST containers too.
 
 ## Pattern 1: Simple mock override
 

modern_di/container.py

@@ -9,7 +9,7 @@
 from modern_di.providers.container_provider import container_provider
 from modern_di.registries.cache_registry import CacheRegistry
 from modern_di.registries.context_registry import ContextRegistry
-from modern_di.registries.overrides_registry import OverridesRegistry
+from modern_di.registries.overrides_registry import OverrideHandle, OverridesRegistry
 from modern_di.registries.providers_registry import ProvidersRegistry
 from modern_di.scope import Scope
 
@@ -338,8 +338,19 @@ def close_sync(self) -> None:
         finally:
             self.closed = True
 
-    def override(self, provider: AbstractProvider[types.T], override_object: types.T) -> None:
+    def override(self, provider: AbstractProvider[types.T], override_object: types.T) -> OverrideHandle[types.T]:
+        """Apply an override immediately.
+
+        Use the returned handle as a context manager to auto-restore the prior state.
+        """
+        prior = self.overrides_registry.fetch_override(provider.provider_id)
         self.overrides_registry.override(provider.provider_id, override_object)
+        return OverrideHandle(
+            registry=self.overrides_registry,
+            provider_id=provider.provider_id,
+            prior=prior,
+            override_object=override_object,
+        )
 
     def reset_override(self, provider: AbstractProvider[types.T] | None = None) -> None:
         self.overrides_registry.reset_override(provider.provider_id if provider else None)

modern_di/registries/overrides_registry.py

@@ -1,5 +1,6 @@
 import dataclasses
 import typing
+from types import TracebackType
 
 from modern_di import types
 
@@ -19,3 +20,40 @@ def reset_override(self, provider_id: int | None = None) -> None:
 
     def fetch_override(self, provider_id: int) -> object:
         return self.overrides.get(provider_id, types.UNSET)
+
+
+class OverrideHandle(typing.Generic[types.T]):
+    """Context-manager handle returned by ``Container.override``.
+
+    The override is already active when the handle is created; ``__exit__`` restores the
+    snapshot taken at creation — the prior override, or no override. Single-use contract.
+    """
+
+    __slots__ = ("_prior", "_provider_id", "_registry", "override_object")
+
+    def __init__(
+        self,
+        *,
+        registry: OverridesRegistry,
+        provider_id: int,
+        prior: object,
+        override_object: types.T,
+    ) -> None:
+        self._registry = registry
+        self._provider_id = provider_id
+        self._prior = prior
+        self.override_object = override_object
+
+    def __enter__(self) -> types.T:
+        return self.override_object
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        if isinstance(self._prior, types.UnsetType):
+            self._registry.reset_override(self._provider_id)
+        else:
+            self._registry.override(self._provider_id, self._prior)

planning/changes/2026-07-08.03-override-context-manager.md

@@ -0,0 +1,77 @@
+---
+summary: INT-4 — Container.override returns an OverrideHandle usable as a context manager; __exit__ restores the provider's prior override state (stacking-safe), imperative callers unaffected.
+---
+
+# Design: Self-resetting context-manager form of override
+
+## Summary
+
+The accepted INT-4 item from the 2026-07-05 3.0 UX research.
+`container.override(provider, obj)` keeps applying the override immediately
+and now returns a small `OverrideHandle`; used as a context manager, its
+`__exit__` restores the provider's prior override state — the previously
+stacked override if there was one, otherwise no override. `__enter__`
+returns the override object, so `with container.override(p, mock) as m:`
+binds the mock. Imperative callers that ignore the return value see
+identical behavior.
+
+## Motivation
+
+Research item INT-4 (verified): `override()`/`reset_override()` is
+imperative-only, so tests must pair the calls manually — the
+forget-to-reset leak FastAPI's `dependency_overrides` docs warn about. Both
+field frameworks with runtime overrides treat auto-reset as the primary
+testing spelling: dependency-injector's
+`with container.api_client_factory.override(Mock(...)):` and wireup's
+`with container.override.injectable(...)` (nesting supported since 2.8).
+
+## Design
+
+- `OverrideHandle` — a slotted class in
+  `modern_di/registries/overrides_registry.py`, constructed by
+  `Container.override` with the registry, the provider id, and a snapshot
+  of the prior value (`fetch_override` result: the previous override or
+  `UNSET`) taken before the new override is written.
+- `__enter__` returns the override object. `__exit__` restores the
+  snapshot: re-set the prior override, or pop the entry when the snapshot
+  is `UNSET`. Restoration is unconditional — deterministic even if
+  `reset_override()` or further `override()` calls ran inside the block.
+- Nesting: each `with` level snapshots the level above, so unwinding
+  restores outer overrides in order (the registry stays a flat dict; the
+  stack lives in the handles).
+- `Container.override` return type changes `None` → `OverrideHandle`; the
+  override is active from the call, not from `__enter__` (imperative
+  compatibility). `reset_override` is untouched.
+- Single-use is the documented contract; the handle holds a snapshot, not
+  live state, so a stray second use merely re-restores the same snapshot.
+- No new exceptions, no warnings, no changes to resolution or the
+  registry's dict semantics.
+
+## Non-goals
+
+- No separate `overriding()` method — one spelling, per the ruling.
+- No provider-level `.override()` (dependency-injector style) — overrides
+  stay a Container concern.
+- No fixture-scoped pytest integration here — modern-di-pytest (sibling
+  repo) builds on the handle.
+
+## Testing
+
+TDD; `just test-ci` 100% line coverage. Cases: `with` applies and
+auto-resets; `as` binds the override object; prior imperative override is
+restored after the block; nested same-provider overrides unwind to the
+outer mock then to none; an exception inside the block still restores;
+`reset_override()` inside the block then exit restores the snapshot;
+imperative call ignoring the handle behaves exactly as before (existing
+tests untouched). `just lint-ci`, `just docs-build` strict,
+`just check-planning` green.
+
+## Risk
+
+- Return-type change surprises a caller asserting `is None` — no such
+  caller exists in-repo or in documented patterns; the research verified
+  the change is invisible to callers that ignore it. Low.
+- Snapshot semantics vs "clear on exit" intuition when code inside the
+  block mutates overrides — mitigated: unconditional snapshot restore is
+  the simplest deterministic rule and is documented in the recipes page.
+  Low.

tests/test_container.py

@@ -938,3 +938,63 @@ class G(Group):
 
     assert request_container.resolve_dependency(str) == "value"
     assert request_container.resolve_dependency(G.request_factory) == "value"
+
+
+class _OverrideSvc: ...
+
+
+class _OverrideGroup(Group):
+    svc = providers.Factory(_OverrideSvc)
+
+
+def test_override_context_manager_applies_and_resets() -> None:
+    container = Container(groups=[_OverrideGroup], validate=False)
+    mock = _OverrideSvc()
+    with container.override(_OverrideGroup.svc, mock) as bound:
+        assert bound is mock
+        assert container.resolve(_OverrideSvc) is mock
+    assert container.resolve(_OverrideSvc) is not mock
+
+
+def test_override_context_manager_restores_prior_imperative_override() -> None:
+    container = Container(groups=[_OverrideGroup], validate=False)
+    first = _OverrideSvc()
+    second = _OverrideSvc()
+    container.override(_OverrideGroup.svc, first)
+    with container.override(_OverrideGroup.svc, second):
+        assert container.resolve(_OverrideSvc) is second
+    assert container.resolve(_OverrideSvc) is first
+
+
+def test_override_context_manager_nested_unwinds_in_order() -> None:
+    container = Container(groups=[_OverrideGroup], validate=False)
+    outer = _OverrideSvc()
+    inner = _OverrideSvc()
+    with container.override(_OverrideGroup.svc, outer):
+        with container.override(_OverrideGroup.svc, inner):
+            assert container.resolve(_OverrideSvc) is inner
+        assert container.resolve(_OverrideSvc) is outer
+    resolved = container.resolve(_OverrideSvc)
+    assert resolved is not outer
+    assert resolved is not inner
+
+
+def test_override_context_manager_restores_on_exception() -> None:
+    container = Container(groups=[_OverrideGroup], validate=False)
+    mock = _OverrideSvc()
+    msg = "boom"
+    with pytest.raises(RuntimeError), container.override(_OverrideGroup.svc, mock):
+        raise RuntimeError(msg)
+    assert container.resolve(_OverrideSvc) is not mock
+
+
+def test_override_context_manager_exit_restores_snapshot_after_inner_reset() -> None:
+    container = Container(groups=[_OverrideGroup], validate=False)
+    first = _OverrideSvc()
+    second = _OverrideSvc()
+    container.override(_OverrideGroup.svc, first)
+    with container.override(_OverrideGroup.svc, second):
+        container.reset_override(_OverrideGroup.svc)
+        assert container.resolve(_OverrideSvc) is not first
+        assert container.resolve(_OverrideSvc) is not second
+    assert container.resolve(_OverrideSvc) is first  # exit restores the snapshot taken at override() time