Coverage for src/lexigram/web/background/tasks.py: 46%

1from __future__ import annotations 

2 

3import asyncio 

4import contextvars 

5from typing import TYPE_CHECKING, Any 

6 

7from lexigram.contracts.web.protocols import BackgroundTaskRunnerProtocol 

8from lexigram.primitives.context import propagate_context, with_context 

9 

10if TYPE_CHECKING: 

11 from collections.abc import Callable 

12 

13 

14def _is_async_callable(obj: Any) -> bool: 

15 """Return ``True`` for async functions *and* async callable objects. 

16 

17 ``asyncio.iscoroutinefunction`` returns ``False`` for instances whose 

18 class defines ``async def __call__``. This helper covers both cases so 

19 that callable objects with an async ``__call__`` are awaited correctly. 

20 """ 

21 return asyncio.iscoroutinefunction(obj) or asyncio.iscoroutinefunction( 

22 getattr(type(obj), "__call__", None) 

23 ) 

24 

25 

26class BackgroundTasks: 

27 """Accumulates background tasks to run after response is sent. 

28 

29 Context is captured at :meth:`add` time using :func:`propagate_context` 

30 so that :meth:`_execute_all` restores the request-scope context variables 

31 (request_id, trace_id, …) even when called after the original request 

32 scope has been torn down. 

33 """ 

34 

35 def __init__(self) -> None: 

36 self._tasks: list[ 

37 tuple[ 

38 Callable[..., Any], tuple[Any, ...], dict[str, Any], contextvars.Context 

39 ] 

40 ] = [] 

41 

42 def add(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> None: 

43 """Add a task to be executed after the response. 

44 

45 The current :mod:`contextvars` context is snapshotted immediately so 

46 that :meth:`_execute_all` can restore it at execution time. 

47 """ 

48 ctx = propagate_context() 

49 self._tasks.append((func, args, kwargs, ctx)) 

50 

51 async def _execute_all(self) -> None: 

52 """Execute all queued tasks, each with its captured context restored.""" 

53 for func, args, kwargs, ctx in self._tasks: 

54 with with_context(ctx): 

55 if _is_async_callable(func): 

56 await func(*args, **kwargs) 

57 else: 

58 func(*args, **kwargs) 

59 

60 def __len__(self) -> int: 

61 return len(self._tasks) 

62 

63 def __bool__(self) -> bool: 

64 return bool(self._tasks) 

65 

66 

67class StarletteBackgroundTaskRunner(BackgroundTaskRunnerProtocol): 

68 """Starlette-backed implementation of BackgroundTaskRunnerProtocol. 

69 

70 Wraps Starlette's BackgroundTasks to provide the framework's 

71 BackgroundTaskRunnerProtocol without leaking Starlette types 

72 into the web layer's public API surface. 

73 

74 Example: 

75 ```python 

76 runner = StarletteBackgroundTaskRunner() 

77 runner.add_task(send_email, to="user@example.com") 

78 response = JSONResponse(data, background=runner._to_starlette()) 

79 ``` 

80 """ 

81 

82 def __init__(self) -> None: 

83 from starlette.background import BackgroundTasks as _StarletteBackgroundTasks 

84 

85 self._starlette_tasks = _StarletteBackgroundTasks() 

86 

87 def add_task(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> None: 

88 """Add a function to run in the background after the response is sent. 

89 

90 The current :mod:`contextvars` context is snapshotted immediately so 

91 that the task sees request-scope variables (request_id, trace_id, …) 

92 even if the originating request scope has been torn down by the time 

93 Starlette executes the task. 

94 

95 Sync callables are wrapped with a sync closure so that Starlette still 

96 dispatches them to its threadpool via ``run_in_executor``. Async 

97 callables (including instances with ``async def __call__``) are wrapped 

98 with an async closure so they are awaited directly on the event loop. 

99 

100 Args: 

101 func: Async or sync callable to execute. 

102 *args: Positional arguments for the callable. 

103 **kwargs: Keyword arguments for the callable. 

104 """ 

105 ctx = propagate_context() 

106 

107 if _is_async_callable(func): 

108 

109 async def _async_wrapper(*a: Any, **kw: Any) -> None: 

110 with with_context(ctx): 

111 await func(*a, **kw) 

112 

113 self._starlette_tasks.add_task(_async_wrapper, *args, **kwargs) 

114 else: 

115 

116 def _sync_wrapper(*a: Any, **kw: Any) -> None: 

117 with with_context(ctx): 

118 func(*a, **kw) 

119 

120 self._starlette_tasks.add_task(_sync_wrapper, *args, **kwargs) 

121 

122 def _to_starlette(self) -> Any: 

123 """Return the underlying Starlette BackgroundTasks instance. 

124 

125 This is an internal method for use by the web routing layer only. 

126 It must not appear in the public API. 

127 """ 

128 return self._starlette_tasks 

129 

130 

131class BackgroundTaskScope: 

132 """Dependency injection scope for BackgroundTasks.""" 

133 

134 def __init__(self) -> None: 

135 self._tasks = BackgroundTasks() 

136 

137 @property 

138 def tasks(self) -> BackgroundTasks: 

139 return self._tasks 

140 

141 def add(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> None: 

142 """Add a background task.""" 

143 self._tasks.add(func, *args, **kwargs) 

144 

145 

146__all__ = [ 

147 "BackgroundTaskScope", 

148 "BackgroundTasks", 

149 "StarletteBackgroundTaskRunner", 

150]