16081423623

Committed 04 Jul 2025 09:57PM UTC coverage: 97.685%. First build

Build # 16081423623

Build Type

Pull #33

github

Committed by

web-flow

Commit Message

Merge 6eb57167c into 784a68391

Pull Request Pull Request #33: Add Subscriptions

Run Details

1798 of 1841 branches covered (97.66%)

Branch coverage included in aggregate %.

1009 of 1176 new or added lines in 36 files covered. (85.8%)

26853 of 27489 relevant lines covered (97.69%)

8.79 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

76.79

/undine/hooks.py

from __future__ import annotations

import dataclasses
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Awaitable, Callable, Generator
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from functools import wraps
from typing import TYPE_CHECKING, Any, Self, TypeVar

if TYPE_CHECKING:
    from collections.abc import AsyncIterator
    from types import TracebackType

    from graphql import DocumentNode, ExecutionResult

    from undine.dataclasses import GraphQLHttpParams
    from undine.typing import DjangoRequestProtocol

__all__ = [
    "LifecycleHook",
    "LifecycleHookContext",
    "LifecycleHookManager",
    "delegate_to_subgenerator",
    "use_lifecycle_hooks_async",
    "use_lifecycle_hooks_sync",
]


@dataclasses.dataclass(slots=True, kw_only=True)
class LifecycleHookContext:
    """Context passed to a lifecycle hook."""

    source: str
    """Source GraphQL document string."""

    document: DocumentNode | None
    """Parsed GraphQL document AST. Available after parsing is complete."""

    variables: dict[str, Any]
    """Variables passed to the GraphQL operation."""

    operation_name: str | None
    """Name of the GraphQL operation."""

    extensions: dict[str, Any]
    """GraphQL operation extensions received from the client."""

    request: DjangoRequestProtocol
    """Django request during which the GraphQL request is being executed."""

    result: ExecutionResult | Awaitable[ExecutionResult | AsyncIterator[ExecutionResult]] | None
    """Execution result of the GraphQL operation. Adding a result here will cause an early exit."""

    @classmethod
    def from_graphql_params(cls, params: GraphQLHttpParams, request: DjangoRequestProtocol) -> Self:
        return cls(
            source=params.document,
            document=None,
            variables=params.variables,
            operation_name=params.operation_name,
            extensions=params.extensions,
            request=request,
            result=None,
        )


@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
class LifecycleHook(ABC):
    """Base class for lifecycle hooks."""

    context: LifecycleHookContext

    @contextmanager
    def use_sync(self) -> Generator[None, None, None]:
        yield from self.run()

    @asynccontextmanager
    async def use_async(self) -> AsyncGenerator[None, None]:
        gen = self.run_async()
        async with delegate_to_subgenerator(gen):
            async for _ in gen:
                yield

    @abstractmethod
    def run(self) -> Generator[None, None, None]:
        """
        Override this method to define how the hook should be executed.
        Anything before the yield statement will be executed before the hooking point.
        Anything after the yield statement will be executed after the hooking point.
        """
        yield

    async def run_async(self) -> AsyncGenerator[None, None]:
        """
        Override this method to define how the hook should be executed in an async context.
        Uses the `run` method by default.
        """
        with delegate_to_subgenerator(self.run()) as gen:
            for _ in gen:
                yield


TLifecycleHook = TypeVar("TLifecycleHook", bound=LifecycleHook)


class LifecycleHookManager(ExitStack, AsyncExitStack):
    """Allows executing multiple lifecycle hooks at once."""

    def __init__(self, *, hooks: list[type[TLifecycleHook]], context: LifecycleHookContext) -> None:
        self.hooks: list[TLifecycleHook] = [hook(context=context) for hook in hooks]
        super().__init__()

    def __enter__(self) -> Self:
        for hook in self.hooks:
            self.enter_context(hook.use_sync())
        return super().__enter__()

    async def __aenter__(self) -> Self:
        for hook in self.hooks:
            await self.enter_async_context(hook.use_async())
        return await super().__aenter__()


R = TypeVar("R")
HookableSync = Callable[[LifecycleHookContext], R]
HookableAsync = Callable[[LifecycleHookContext], Awaitable[R]]


def use_lifecycle_hooks_sync(hooks: list[type[TLifecycleHook]]) -> Callable[[HookableSync[R]], HookableSync[R]]:
    """Run given function using the given lifecycle hooks."""

    def decorator(func: HookableSync[R]) -> HookableSync[R]:
        @wraps(func)
        def wrapper(context: LifecycleHookContext) -> R:
            with LifecycleHookManager(hooks=hooks, context=context):
                return func(context)

        return wrapper

    return decorator


def use_lifecycle_hooks_async(hooks: list[type[TLifecycleHook]]) -> Callable[[HookableAsync[R]], HookableAsync[R]]:
    """Run given function using the given lifecycle hooks."""

    def decorator(func: HookableAsync[R]) -> HookableAsync[R]:
        @wraps(func)
        async def wrapper(context: LifecycleHookContext) -> R:  # type: ignore[return]
            async with LifecycleHookManager(hooks=hooks, context=context):
                return await func(context)

        return wrapper

    return decorator


class delegate_to_subgenerator:  # noqa: N801
    """
    Allows delegating how a generator exists to a subgenerator.

    >>> def subgenerator():
    ...     for _ in range(2):
    ...         yield
    >>>
    >>> def generator():
    >>>     with delegate_to_subgenerator(subgenerator()) as sub:
    ...         for _ in sub:
    ...             yield
    >>>
    >>> for item in generator():
    ...     pass

    If the generator exists normally, the subgenerator will be closed.
    If the generator exists with an exception, the error is propagated to the subgenerator
    so that it may handle the error.
    """

    def __init__(self, gen: Generator[None, None, None] | AsyncGenerator[None, None]) -> None:
        """
        Allows delegating how a generator exists to a subgenerator.

        :param gen: The generator to delegate to. If generator is an async generator,
                    must use `async with` syntax to delegate. For regular generators,
                    plain `with` syntax must be used.
        """
        self.gen = gen

    def __enter__(self) -> Generator[None, None, None]:
        if not isinstance(self.gen, Generator):
            msg = "Given object is not a Generator"
            raise TypeError(msg)

        return self.gen

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool:
        if not isinstance(self.gen, Generator):  # type: ignore[unreachable]
            msg = "Given object is not a Generator"
            raise TypeError(msg)

        # If no exception was raised, close the generator.
        if exc_type is None:
            self.gen.close()
            return False

        # Otherwise, allow the subgenerator to handle the exception.
        # This has mostly been copied from `contextlib._GeneratorContextManager.__exit__`.
        if exc_value is None:
            exc_value = exc_type()

        try:
            self.gen.throw(exc_value)

        except StopIteration as error:
            return error is not exc_value

        except RuntimeError as error:
            if error is exc_value:
                error.__traceback__ = traceback
                return False
            if isinstance(exc_value, StopIteration) and error.__cause__ is exc_value:
                exc_value.__traceback__ = traceback
                return False
            raise

        except BaseException as error:
            if error is not exc_value:
                raise
            error.__traceback__ = traceback
            return False

        try:
            msg = "generator didn't stop after throw()"
            raise RuntimeError(msg)
        finally:
            self.gen.close()

    async def __aenter__(self) -> AsyncGenerator[None, None]:
        if not isinstance(self.gen, AsyncGenerator):
            msg = "Given object is not an AsyncGenerator"
            raise TypeError(msg)

        return self.gen

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> bool:
        if not isinstance(self.gen, AsyncGenerator):
            msg = "Given object is not an AsyncGenerator"
            raise TypeError(msg)

        # If no exception was raised, close the generator.
        if exc_type is None:
            await self.gen.aclose()
            return False

        # Otherwise, allow the subgenerator to handle the exception.
        # This has mostly been copied from `contextlib._AsyncGeneratorContextManager.__aexit__`.
        if exc_value is None:
            exc_value = exc_type()

        try:
            await self.gen.athrow(exc_value)

        except StopAsyncIteration as error:
            return error is not exc_value

        except RuntimeError as error:
            if error is exc_value:
                error.__traceback__ = traceback
                return False
            if isinstance(exc_value, (StopIteration, StopAsyncIteration)) and error.__cause__ is exc_value:
                exc_value.__traceback__ = traceback
                return False
            raise

        except BaseException as error:
            if error is not exc_value:
                raise
            error.__traceback__ = traceback
            return False

        try:
            msg = "generator didn't stop after athrow()"
            raise RuntimeError(msg)
        finally:
            await self.gen.aclose()

1	from __future__ import annotations	9✔
2
3	import dataclasses	9✔
4	from abc import ABC, abstractmethod	9✔
5	from collections.abc import AsyncGenerator, Awaitable, Callable, Generator	9✔
6	from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager	9✔
7	from functools import wraps	9✔
8	from typing import TYPE_CHECKING, Any, Self, TypeVar	9✔
9
10	if TYPE_CHECKING:
11	from collections.abc import AsyncIterator
12	from types import TracebackType
13
14	from graphql import DocumentNode, ExecutionResult
15
16	from undine.dataclasses import GraphQLHttpParams
17	from undine.typing import DjangoRequestProtocol
18
19	__all__ = [	9✔
20	"LifecycleHook",
21	"LifecycleHookContext",
22	"LifecycleHookManager",
23	"delegate_to_subgenerator",
24	"use_lifecycle_hooks_async",
25	"use_lifecycle_hooks_sync",
26	]
27
28
29	@dataclasses.dataclass(slots=True, kw_only=True)	9✔
30	class LifecycleHookContext:	9✔
31	"""Context passed to a lifecycle hook."""
32
33	source: str	9✔
34	"""Source GraphQL document string."""	9✔
35
36	document: DocumentNode \| None	9✔
37	"""Parsed GraphQL document AST. Available after parsing is complete."""	9✔
38
39	variables: dict[str, Any]	9✔
40	"""Variables passed to the GraphQL operation."""
41
42	operation_name: str \| None	9✔
43	"""Name of the GraphQL operation."""	9✔
44
45	extensions: dict[str, Any]	9✔
46	"""GraphQL operation extensions received from the client."""	9✔
47
48	request: DjangoRequestProtocol	9✔
49	"""Django request during which the GraphQL request is being executed."""	9✔
50
51	result: ExecutionResult \| Awaitable[ExecutionResult \| AsyncIterator[ExecutionResult]] \| None	9✔
52	"""Execution result of the GraphQL operation. Adding a result here will cause an early exit."""	9✔
53
54	@classmethod	9✔
55	def from_graphql_params(cls, params: GraphQLHttpParams, request: DjangoRequestProtocol) -> Self:	9✔
56	return cls(	9✔
57	source=params.document,
58	document=None,
59	variables=params.variables,
60	operation_name=params.operation_name,
61	extensions=params.extensions,
62	request=request,
63	result=None,
64	)
65
66
67	@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)	9✔
68	class LifecycleHook(ABC):	9✔
69	"""Base class for lifecycle hooks."""
70
71	context: LifecycleHookContext	9✔
72
73	@contextmanager	9✔
74	def use_sync(self) -> Generator[None, None, None]:	9✔
75	yield from self.run()	9✔
76
77	@asynccontextmanager	9✔
78	async def use_async(self) -> AsyncGenerator[None, None]:	9✔
79	gen = self.run_async()	9✔
80	async with delegate_to_subgenerator(gen):	9✔
81	async for _ in gen:	9✔
82	yield	9✔
83
84	@abstractmethod	9✔
85	def run(self) -> Generator[None, None, None]:	9✔
86	"""
87	Override this method to define how the hook should be executed.
88	Anything before the yield statement will be executed before the hooking point.
89	Anything after the yield statement will be executed after the hooking point.
90	"""
91	yield	×
92
93	async def run_async(self) -> AsyncGenerator[None, None]:	9✔
94	"""
95	Override this method to define how the hook should be executed in an async context.
96	Uses the `run` method by default.
97	"""
98	with delegate_to_subgenerator(self.run()) as gen:	9✔
99	for _ in gen:	9✔
100	yield	9✔
101
102
103	TLifecycleHook = TypeVar("TLifecycleHook", bound=LifecycleHook)	9✔
104
105
106	class LifecycleHookManager(ExitStack, AsyncExitStack):	9✔
107	"""Allows executing multiple lifecycle hooks at once."""
108
109	def __init__(self, *, hooks: list[type[TLifecycleHook]], context: LifecycleHookContext) -> None:	9✔
110	self.hooks: list[TLifecycleHook] = [hook(context=context) for hook in hooks]	9✔
111	super().__init__()	9✔
112
113	def __enter__(self) -> Self:	9✔
114	for hook in self.hooks:	9✔
115	self.enter_context(hook.use_sync())	9✔
116	return super().__enter__()	9✔
117
118	async def __aenter__(self) -> Self:	9✔
119	for hook in self.hooks:	9✔
120	await self.enter_async_context(hook.use_async())	9✔
121	return await super().__aenter__()	9✔
122
123
124	R = TypeVar("R")	9✔
125	HookableSync = Callable[[LifecycleHookContext], R]	9✔
126	HookableAsync = Callable[[LifecycleHookContext], Awaitable[R]]	9✔
127
128
129	def use_lifecycle_hooks_sync(hooks: list[type[TLifecycleHook]]) -> Callable[[HookableSync[R]], HookableSync[R]]:	9✔
130	"""Run given function using the given lifecycle hooks."""
131
132	def decorator(func: HookableSync[R]) -> HookableSync[R]:	9✔
133	@wraps(func)	9✔
134	def wrapper(context: LifecycleHookContext) -> R:	9✔
135	with LifecycleHookManager(hooks=hooks, context=context):	9✔
136	return func(context)	9✔
137
138	return wrapper	9✔
139
140	return decorator	9✔
141
142
143	def use_lifecycle_hooks_async(hooks: list[type[TLifecycleHook]]) -> Callable[[HookableAsync[R]], HookableAsync[R]]:	9✔
144	"""Run given function using the given lifecycle hooks."""
145
146	def decorator(func: HookableAsync[R]) -> HookableAsync[R]:	9✔
147	@wraps(func)	9✔
148	async def wrapper(context: LifecycleHookContext) -> R: # type: ignore[return]	9✔
149	async with LifecycleHookManager(hooks=hooks, context=context):	9✔
150	return await func(context)	9✔
151
152	return wrapper	9✔
153
154	return decorator	9✔
155
156
157	class delegate_to_subgenerator: # noqa: N801	9✔
158	"""
159	Allows delegating how a generator exists to a subgenerator.
160
161	>>> def subgenerator():
162	... for _ in range(2):
163	... yield
164	>>>
165	>>> def generator():
166	>>> with delegate_to_subgenerator(subgenerator()) as sub:
167	... for _ in sub:
168	... yield
169	>>>
170	>>> for item in generator():
171	... pass
172
173	If the generator exists normally, the subgenerator will be closed.
174	If the generator exists with an exception, the error is propagated to the subgenerator
175	so that it may handle the error.
176	"""
177
178	def __init__(self, gen: Generator[None, None, None] \| AsyncGenerator[None, None]) -> None:	9✔
179	"""
180	Allows delegating how a generator exists to a subgenerator.
181
182	:param gen: The generator to delegate to. If generator is an async generator,
183	must use `async with` syntax to delegate. For regular generators,
184	plain `with` syntax must be used.
185	"""
186	self.gen = gen	9✔
187
188	def __enter__(self) -> Generator[None, None, None]:	9✔
189	if not isinstance(self.gen, Generator):	9✔
NEW 190	msg = "Given object is not a Generator"	×
NEW 191	raise TypeError(msg)	×
192
193	return self.gen	9✔
194
195	def __exit__(	9✔
196	self,
197	exc_type: type[BaseException] \| None,
198	exc_value: BaseException \| None,
199	traceback: TracebackType \| None,
200	) -> bool:
201	if not isinstance(self.gen, Generator): # type: ignore[unreachable]	9✔
NEW 202	msg = "Given object is not a Generator"	×
NEW 203	raise TypeError(msg)	×
204
205	# If no exception was raised, close the generator.
206	if exc_type is None:	9✔
207	self.gen.close()	9✔
208	return False	9✔
209
210	# Otherwise, allow the subgenerator to handle the exception.
211	# This has mostly been copied from `contextlib._GeneratorContextManager.__exit__`.
212	if exc_value is None:	9✔
NEW 213	exc_value = exc_type()	×
214
215	try:	9✔
216	self.gen.throw(exc_value)	9✔
217
218	except StopIteration as error:	9✔
219	return error is not exc_value	9✔
220
221	except RuntimeError as error:	9✔
NEW 222	if error is exc_value:	×
NEW 223	error.__traceback__ = traceback	×
NEW 224	return False	×
NEW 225	if isinstance(exc_value, StopIteration) and error.__cause__ is exc_value:	×
NEW 226	exc_value.__traceback__ = traceback	×
NEW 227	return False	×
NEW 228	raise	×
229
230	except BaseException as error:	9✔
231	if error is not exc_value:	9✔
NEW 232	raise	×
233	error.__traceback__ = traceback	9✔
234	return False	9✔
235
NEW 236	try:	×
NEW 237	msg = "generator didn't stop after throw()"	×
NEW 238	raise RuntimeError(msg)	×
239	finally:
NEW 240	self.gen.close()	×
241
242	async def __aenter__(self) -> AsyncGenerator[None, None]:	9✔
243	if not isinstance(self.gen, AsyncGenerator):	9✔
NEW 244	msg = "Given object is not an AsyncGenerator"	×
NEW 245	raise TypeError(msg)	×
246
247	return self.gen	9✔
248
249	async def __aexit__(	9✔
250	self,
251	exc_type: type[BaseException] \| None,
252	exc_value: BaseException \| None,
253	traceback: TracebackType \| None,
254	) -> bool:
255	if not isinstance(self.gen, AsyncGenerator):	9✔
NEW 256	msg = "Given object is not an AsyncGenerator"	×
NEW 257	raise TypeError(msg)	×
258
259	# If no exception was raised, close the generator.
260	if exc_type is None:	9✔
261	await self.gen.aclose()	9✔
262	return False	9✔
263
264	# Otherwise, allow the subgenerator to handle the exception.
265	# This has mostly been copied from `contextlib._AsyncGeneratorContextManager.__aexit__`.
266	if exc_value is None:	9✔
NEW 267	exc_value = exc_type()	×
268
269	try:	9✔
270	await self.gen.athrow(exc_value)	9✔
271
272	except StopAsyncIteration as error:	9✔
273	return error is not exc_value	9✔
274
275	except RuntimeError as error:	9✔
NEW 276	if error is exc_value:	×
NEW 277	error.__traceback__ = traceback	×
NEW 278	return False	×
NEW 279	if isinstance(exc_value, (StopIteration, StopAsyncIteration)) and error.__cause__ is exc_value:	×
NEW 280	exc_value.__traceback__ = traceback	×
NEW 281	return False	×
NEW 282	raise	×
283
284	except BaseException as error:	9✔
285	if error is not exc_value:	9✔
NEW 286	raise	×
287	error.__traceback__ = traceback	9✔
288	return False	9✔
289
NEW 290	try:	×
NEW 291	msg = "generator didn't stop after athrow()"	×
NEW 292	raise RuntimeError(msg)	×
293	finally:
NEW 294	await self.gen.aclose()	×

MrThearMan / undine / 16081423623

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous