Merge pull request #65 from kraken-tech/meshy/contextmanager-only-decorator

Rework how `savepoint` avoids use as a decorator

Author: meshy

CHANGELOG.md

@@ -13,10 +13,8 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ### Removed
 
-- `django_subatomic.db.NotADecorator` is no longer a part of the public API.
-  It has been renamed to `_NotADecorator`.
-  This implementation detail was not intended for general use,
-  and may be removed in a future release.
+- `_contextmanager_without_decorator`, `_NonDecoratorContextManager`, `NotADecorator`.
+  These implementation details were not intended to be a part of our public API.
 
 ## [0.1.1] - 2025-09-20
 

src/django_subatomic/_utils.py

@@ -0,0 +1,80 @@
+import functools
+from collections.abc import Callable, Generator
+from contextlib import AbstractContextManager
+from types import TracebackType
+from typing import Literal
+
+
+def contextmanager[**P_Args, T_YieldType](
+    func: Callable[P_Args, Generator[T_YieldType]],
+) -> Callable[P_Args, AbstractContextManager[T_YieldType]]:
+    """
+    Make a generator function into a context manager.
+
+    Similar to `contextlib.contextmanager`, except it does not allow the
+    returned function to be used as a decorator
+    """
+
+    @functools.wraps(func)
+    def wrapper(
+        *args: P_Args.args, **kwargs: P_Args.kwargs
+    ) -> AbstractContextManager[T_YieldType]:
+        gen = func(*args, **kwargs)
+        return _ContextManagerOnly(gen)
+
+    return wrapper
+
+
+class DidNotYield(Exception):
+    pass
+
+
+class UnexpectedSecondYield(Exception):
+    pass
+
+
+class _ContextManagerOnly[T_YieldType]:
+    """
+    Wraps a generator to act as a context manager.
+    """
+
+    def __init__(self, gen: Generator[T_YieldType]) -> None:
+        self.gen = gen
+
+    def __enter__(self) -> T_YieldType:
+        try:
+            # Run the generator up until the 'yield'.
+            return next(self.gen)
+        except StopIteration as e:
+            raise DidNotYield from e
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> Literal[True] | None:
+        # An exception was raised from the 'with' block.
+        if exc_val is not None:
+            try:
+                # Throw the exception into the generator's 'yield' point.
+                self.gen.throw(exc_val)
+            except StopIteration:
+                # The generator handled the exception. Returning `True`
+                # suppresses the exception from propagating.
+                return True
+            else:
+                # The generator handled the exception but then yielded again.
+                raise UnexpectedSecondYield
+
+        # The 'with' block completed without exception.
+        else:
+            try:
+                # Run the generator after the 'yield'.
+                next(self.gen)
+            except StopIteration:
+                # A clean exit.
+                return None
+            else:
+                # The generator yielded a second time.
+                raise UnexpectedSecondYield

src/django_subatomic/db.py

@@ -9,10 +9,11 @@
 from django.conf import settings
 from django.db import transaction as django_transaction
 
+from . import _utils
+
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator, Iterator
-    from typing import NoReturn
 
 
 @contextlib.contextmanager
@@ -63,44 +64,7 @@ def transaction_if_not_already(*, using: str | None = None) -> Iterator[None]:
         yield
 
 
-class _NotADecorator(Exception):
-    """
-    Raised when a context manager is mistakenly used as a decorator.
-    """
-
-
-class _NonDecoratorContextManager[T_Co](contextlib._GeneratorContextManager[T_Co]):  # noqa: SLF001
-    """
-    Hacked version of contextlib._GeneratorContextManager that prevents use as a decorator.
-
-    Use _contextmanager_without_decorator to create instances of this class.
-
-    This is a weird hack, but it beats copying the bulk of the
-    contextlib.contextmanager code into our own codebase.
-    """
-
-    def __call__(self, func: object) -> NoReturn:
-        raise _NotADecorator
-
-
-def _contextmanager_without_decorator[**P, T_Co](
-    func: Callable[P, Generator[T_Co, None, None]], /
-) -> Callable[P, _NonDecoratorContextManager[T_Co]]:
-    """
-    Decorate a generator function to make it a context manager.
-
-    This is pretty much the same as `contextlib.contextmanager`,
-    but it prevents use as a decorator.
-    """
-
-    @functools.wraps(func)
-    def helper(*args: P.args, **kwds: P.kwargs) -> _NonDecoratorContextManager[T_Co]:
-        return _NonDecoratorContextManager(func, args, kwds)
-
-    return helper
-
-
-@_contextmanager_without_decorator
+@_utils.contextmanager
 def savepoint(*, using: str | None = None) -> Generator[None, None, None]:
     """
     Create a database savepoint.

tests/test_db.py

@@ -261,9 +261,10 @@ def test_not_a_decorator(self) -> None:
         """
         `savepoint` cannot be used as a decorator.
         """
-        with pytest.raises(db._NotADecorator):  # noqa: SLF001
+        expected_error = "'_ContextManagerOnly' object is not callable"
+        with pytest.raises(TypeError, match=expected_error):
 
-            @db.savepoint()
+            @db.savepoint()  # type: ignore[operator]
             def inner() -> None: ...
 
     @test.part_of_a_transaction()

tests/test_utils.py

@@ -0,0 +1,128 @@
+from collections.abc import Generator
+
+import pytest
+
+from django_subatomic import _utils as utils
+
+
+class _AnError(Exception):
+    pass
+
+
+class TestContextManager:
+    """
+    Tests for the `contextmanager` decorator.
+    """
+
+    def test_context_manager(self) -> None:
+        """
+        Generators that yield once can be used as context managers.
+        """
+
+        @utils.contextmanager
+        def good_example(steps: list[str]) -> Generator[str, None, None]:
+            steps.append("enter")
+            yield "yielded"
+            steps.append("exit")
+
+        steps: list[str] = []
+        with good_example(steps) as yielded:
+            steps.append("body")
+
+        assert steps == ["enter", "body", "exit"]
+        assert yielded == "yielded"
+
+    def test_context_body_raises_exception(self) -> None:
+        """
+        Exceptions raised in the context body are propagated.
+        """
+
+        @utils.contextmanager
+        def basic_context() -> Generator[None, None, None]:
+            yield
+
+        with pytest.raises(_AnError):
+            with basic_context():
+                raise _AnError
+
+    def test_context_manager_handles_exception(self) -> None:
+        """
+        Context managers may handle exceptions raised in the context body.
+        """
+
+        @utils.contextmanager
+        def handle_exception() -> Generator[None, None, None]:
+            # Make sure the error is raised, but don't propagate it.
+            with pytest.raises(_AnError):
+                yield
+
+        with handle_exception():
+            raise _AnError
+
+    def test_context_manager_handles_exception_but_then_yields(self) -> None:
+        """
+        A second yield must not happen after an exception is handled.
+        """
+
+        @utils.contextmanager
+        def handle_exception_but_yield_twice() -> Generator[None, None, None]:
+            try:
+                yield
+            except _AnError:
+                yield
+
+        with pytest.raises(utils.UnexpectedSecondYield):
+            with handle_exception_but_yield_twice():
+                raise _AnError
+
+    def test_fails_without_yield(self) -> None:
+        """
+        Decorated functions must yield once.
+        """
+
+        @utils.contextmanager
+        def yield_not_run() -> Generator[None, None, None]:
+            if False:
+                # Yield makes this a generator, but the yield is never run.
+                yield  # type: ignore[unreachable]
+
+        with pytest.raises(utils.DidNotYield):
+            with yield_not_run():
+                ...
+
+    def test_fails_on_multiple_yields(self) -> None:
+        """
+        Decorated functions must not yield multiple times.
+        """
+
+        @utils.contextmanager
+        def yields_twice(steps: list[str]) -> Generator[None, None, None]:
+            steps.append("enter")
+            yield
+            steps.append("between")
+            yield  # This line causes the failure.
+            # The following line will never be reached.
+            steps.append("after")  # pragma: no cover
+
+        steps: list[str] = []
+        with pytest.raises(utils.UnexpectedSecondYield):
+            with yields_twice(steps):
+                steps.append("body")
+
+        assert steps == ["enter", "body", "between"]
+
+    def test_does_not_work_as_a_decorator(self) -> None:
+        """
+        Functions decorated with `contextmanager` cannot be used as decorators.
+        """
+
+        @utils.contextmanager
+        def not_a_decorator() -> Generator[None, None, None]:
+            yield  # pragma: no cover
+
+        with pytest.raises(
+            TypeError, match="takes 0 positional arguments but 1 was given"
+        ):
+            # This doesn't work as a decorator.
+            @not_a_decorator  # type: ignore[call-arg]
+            def not_reached() -> None: ...