|
16 | 16 | from typing import Any |
17 | 17 |
|
18 | 18 |
|
19 | | -def _sentry_task_factory(loop, coro): |
20 | | - # type: (Any, Any) -> Task[None] |
| 19 | +def patch_asyncio(): |
| 20 | + # type: () -> None |
| 21 | + orig_task_factory = None |
| 22 | + try: |
| 23 | + loop = asyncio.get_running_loop() |
| 24 | + orig_task_factory = loop.get_task_factory() |
21 | 25 |
|
22 | | - async def _coro_creating_hub_and_span(): |
23 | | - # type: () -> None |
24 | | - hub = Hub(Hub.current) |
25 | | - with hub: |
26 | | - with hub.start_span(op=OP.FUNCTION, description=coro.__qualname__): |
27 | | - await coro |
| 26 | + def _sentry_task_factory(loop, coro): |
| 27 | + # type: (Any, Any) -> Any |
28 | 28 |
|
29 | | - # Trying to use user set task factory (if there is one) |
30 | | - orig_factory = loop.get_task_factory() |
31 | | - if orig_factory: |
32 | | - return orig_factory(loop, _coro_creating_hub_and_span) |
| 29 | + async def _coro_creating_hub_and_span(): |
| 30 | + # type: () -> None |
| 31 | + hub = Hub(Hub.current) |
| 32 | + with hub: |
| 33 | + with hub.start_span(op=OP.FUNCTION, description=coro.__qualname__): |
| 34 | + await coro |
33 | 35 |
|
34 | | - # The default task factory in `asyncio` does not have its own function |
35 | | - # but is just a couple of lines in `asyncio.base_events.create_task()` |
36 | | - # Those lines are copied here. |
| 36 | + # Trying to use user set task factory (if there is one) |
| 37 | + if orig_task_factory: |
| 38 | + return orig_task_factory(loop, _coro_creating_hub_and_span()) # type: ignore |
37 | 39 |
|
38 | | - # WARNING: |
39 | | - # If the default behavior of the task creation in asyncio changes, |
40 | | - # this will break! |
41 | | - task = Task(_coro_creating_hub_and_span, loop=loop) # type: ignore |
42 | | - if task._source_traceback: # type: ignore |
43 | | - del task._source_traceback[-1] # type: ignore |
| 40 | + # The default task factory in `asyncio` does not have its own function |
| 41 | + # but is just a couple of lines in `asyncio.base_events.create_task()` |
| 42 | + # Those lines are copied here. |
44 | 43 |
|
45 | | - return task |
| 44 | + # WARNING: |
| 45 | + # If the default behavior of the task creation in asyncio changes, |
| 46 | + # this will break! |
| 47 | + task = Task(_coro_creating_hub_and_span(), loop=loop) |
| 48 | + if task._source_traceback: # type: ignore |
| 49 | + del task._source_traceback[-1] # type: ignore |
46 | 50 |
|
| 51 | + return task |
47 | 52 |
|
48 | | -def patch_asyncio(): |
49 | | - # type: () -> None |
50 | | - try: |
51 | | - loop = asyncio.get_running_loop() |
52 | 53 | loop.set_task_factory(_sentry_task_factory) |
53 | 54 | except RuntimeError: |
54 | 55 | # When there is no running loop, we have nothing to patch. |
|
0 commit comments