Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/asgiref/current_thread_executor.py: 23%

1import threading 

2from collections import deque 

3from collections.abc import Callable 

4from concurrent.futures import Executor, Future 

5from typing import Any, ParamSpec, TypeVar 

6 

7_T = TypeVar("_T") 

8_P = ParamSpec("_P") 

9_R = TypeVar("_R") 

10 

11 

12class _WorkItem: 

13 """ 

14 Represents an item needing to be run in the executor. 

15 Copied from ThreadPoolExecutor (but it's private, so we're not going to rely on importing it) 

16 """ 

17 

18 def __init__( 

19 self, 

20 future: "Future[_R]", 

21 fn: Callable[_P, _R], 

22 *args: _P.args, 

23 **kwargs: _P.kwargs, 

24 ): 

25 self.future = future 

26 self.fn = fn 

27 self.args = args 

28 self.kwargs = kwargs 

29 

30 def run(self) -> None: 

31 __traceback_hide__ = True # noqa: F841 

32 if not self.future.set_running_or_notify_cancel(): 

33 return 

34 try: 

35 result = self.fn(*self.args, **self.kwargs) 

36 except BaseException as exc: 

37 self.future.set_exception(exc) 

38 # Break a reference cycle with the exception 'exc' 

39 self = None # type: ignore[assignment] 

40 else: 

41 self.future.set_result(result) 

42 

43 

44class CurrentThreadExecutor(Executor): 

45 """ 

46 An Executor that actually runs code in the thread it is instantiated in. 

47 Passed to other threads running async code, so they can run sync code in 

48 the thread they came from. 

49 """ 

50 

51 def __init__(self, old_executor: "CurrentThreadExecutor | None") -> None: 

52 self._work_thread = threading.current_thread() 

53 self._work_ready = threading.Condition(threading.Lock()) 

54 self._work_items = deque[_WorkItem]() # synchronized by _work_ready 

55 self._broken = False # synchronized by _work_ready 

56 self._old_executor = old_executor 

57 

58 def run_until_future(self, future: "Future[Any]") -> None: 

59 """ 

60 Runs the code in the work queue until a result is available from the future. 

61 Should be run from the thread the executor is initialised in. 

62 """ 

63 # Check we're in the right thread 

64 if threading.current_thread() != self._work_thread: 

65 raise RuntimeError( 

66 "You cannot run CurrentThreadExecutor from a different thread" 

67 ) 

68 

69 def done(future: "Future[Any]") -> None: 

70 with self._work_ready: 

71 self._broken = True 

72 self._work_ready.notify() 

73 

74 future.add_done_callback(done) 

75 # Keep getting and running work items until the future we're waiting for 

76 # is done and the queue is empty. 

77 while True: 

78 with self._work_ready: 

79 while not self._work_items and not self._broken: 

80 self._work_ready.wait() 

81 if not self._work_items: 

82 break 

83 # Get a work item and run it 

84 work_item = self._work_items.popleft() 

85 work_item.run() 

86 del work_item 

87 

88 def submit( 

89 self, 

90 fn: Callable[_P, _R], 

91 /, 

92 *args: _P.args, 

93 **kwargs: _P.kwargs, 

94 ) -> "Future[_R]": 

95 # Check they're not submitting from the same thread 

96 if threading.current_thread() == self._work_thread: 

97 raise RuntimeError( 

98 "You cannot submit onto CurrentThreadExecutor from its own thread" 

99 ) 

100 f: "Future[_R]" = Future() 

101 work_item = _WorkItem(f, fn, *args, **kwargs) 

102 

103 # Walk up the CurrentThreadExecutor stack to find the closest one still 

104 # running 

105 executor = self 

106 while True: 

107 with executor._work_ready: 

108 if not executor._broken: 

109 # Add to work queue 

110 executor._work_items.append(work_item) 

111 executor._work_ready.notify() 

112 break 

113 if executor._old_executor is None: 

114 raise RuntimeError("CurrentThreadExecutor already quit or is broken") 

115 executor = executor._old_executor 

116 

117 # Return the future 

118 return f