mirror of
				https://github.com/python/cpython.git
				synced 2025-10-31 21:51:50 +00:00 
			
		
		
		
	
		
			
				
	
	
		
			445 lines
		
	
	
	
		
			12 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			445 lines
		
	
	
	
		
			12 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
| import asyncio
 | |
| import io
 | |
| import unittest
 | |
| 
 | |
| 
 | |
| # To prevent a warning "test altered the execution environment"
 | |
| def tearDownModule():
 | |
|     asyncio._set_event_loop_policy(None)
 | |
| 
 | |
| 
 | |
| def capture_test_stack(*, fut=None, depth=1):
 | |
| 
 | |
|     def walk(s):
 | |
|         ret = [
 | |
|             (f"T<{n}>" if '-' not in (n := s.future.get_name()) else 'T<anon>')
 | |
|                 if isinstance(s.future, asyncio.Task) else 'F'
 | |
|         ]
 | |
| 
 | |
|         ret.append(
 | |
|             [
 | |
|                 (
 | |
|                     f"s {entry.frame.f_code.co_name}"
 | |
|                         if entry.frame.f_generator is None else
 | |
|                         (
 | |
|                             f"a {entry.frame.f_generator.cr_code.co_name}"
 | |
|                             if hasattr(entry.frame.f_generator, 'cr_code') else
 | |
|                             f"ag {entry.frame.f_generator.ag_code.co_name}"
 | |
|                         )
 | |
|                 ) for entry in s.call_stack
 | |
|             ]
 | |
|         )
 | |
| 
 | |
|         ret.append(
 | |
|             sorted([
 | |
|                 walk(ab) for ab in s.awaited_by
 | |
|             ], key=lambda entry: entry[0])
 | |
|         )
 | |
| 
 | |
|         return ret
 | |
| 
 | |
|     buf = io.StringIO()
 | |
|     asyncio.print_call_graph(fut, file=buf, depth=depth+1)
 | |
| 
 | |
|     stack = asyncio.capture_call_graph(fut, depth=depth)
 | |
|     return walk(stack), buf.getvalue()
 | |
| 
 | |
| 
 | |
| class CallStackTestBase:
 | |
| 
 | |
|     async def test_stack_tgroup(self):
 | |
| 
 | |
|         stack_for_c5 = None
 | |
| 
 | |
|         def c5():
 | |
|             nonlocal stack_for_c5
 | |
|             stack_for_c5 = capture_test_stack(depth=2)
 | |
| 
 | |
|         async def c4():
 | |
|             await asyncio.sleep(0)
 | |
|             c5()
 | |
| 
 | |
|         async def c3():
 | |
|             await c4()
 | |
| 
 | |
|         async def c2():
 | |
|             await c3()
 | |
| 
 | |
|         async def c1(task):
 | |
|             await task
 | |
| 
 | |
|         async def main():
 | |
|             async with asyncio.TaskGroup() as tg:
 | |
|                 task = tg.create_task(c2(), name="c2_root")
 | |
|                 tg.create_task(c1(task), name="sub_main_1")
 | |
|                 tg.create_task(c1(task), name="sub_main_2")
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_c5[0], [
 | |
|             # task name
 | |
|             'T<c2_root>',
 | |
|             # call stack
 | |
|             ['s c5', 'a c4', 'a c3', 'a c2'],
 | |
|             # awaited by
 | |
|             [
 | |
|                 ['T<anon>',
 | |
|                      ['a _aexit', 'a __aexit__', 'a main', 'a test_stack_tgroup'], []
 | |
|                 ],
 | |
|                 ['T<sub_main_1>',
 | |
|                     ['a c1'],
 | |
|                     [
 | |
|                         ['T<anon>',
 | |
|                             ['a _aexit', 'a __aexit__', 'a main', 'a test_stack_tgroup'], []
 | |
|                         ]
 | |
|                     ]
 | |
|                 ],
 | |
|                 ['T<sub_main_2>',
 | |
|                     ['a c1'],
 | |
|                     [
 | |
|                         ['T<anon>',
 | |
|                             ['a _aexit', 'a __aexit__', 'a main', 'a test_stack_tgroup'], []
 | |
|                         ]
 | |
|                     ]
 | |
|                 ]
 | |
|             ]
 | |
|         ])
 | |
| 
 | |
|         self.assertIn(
 | |
|             ' async CallStackTestBase.test_stack_tgroup()',
 | |
|             stack_for_c5[1])
 | |
| 
 | |
| 
 | |
|     async def test_stack_async_gen(self):
 | |
| 
 | |
|         stack_for_gen_nested_call = None
 | |
| 
 | |
|         async def gen_nested_call():
 | |
|             nonlocal stack_for_gen_nested_call
 | |
|             stack_for_gen_nested_call = capture_test_stack()
 | |
| 
 | |
|         async def gen():
 | |
|             for num in range(2):
 | |
|                 yield num
 | |
|                 if num == 1:
 | |
|                     await gen_nested_call()
 | |
| 
 | |
|         async def main():
 | |
|             async for el in gen():
 | |
|                 pass
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_gen_nested_call[0], [
 | |
|             'T<anon>',
 | |
|             [
 | |
|                 's capture_test_stack',
 | |
|                 'a gen_nested_call',
 | |
|                 'ag gen',
 | |
|                 'a main',
 | |
|                 'a test_stack_async_gen'
 | |
|             ],
 | |
|             []
 | |
|         ])
 | |
| 
 | |
|         self.assertIn(
 | |
|             'async generator CallStackTestBase.test_stack_async_gen.<locals>.gen()',
 | |
|             stack_for_gen_nested_call[1])
 | |
| 
 | |
|     async def test_stack_gather(self):
 | |
| 
 | |
|         stack_for_deep = None
 | |
| 
 | |
|         async def deep():
 | |
|             await asyncio.sleep(0)
 | |
|             nonlocal stack_for_deep
 | |
|             stack_for_deep = capture_test_stack()
 | |
| 
 | |
|         async def c1():
 | |
|             await asyncio.sleep(0)
 | |
|             await deep()
 | |
| 
 | |
|         async def c2():
 | |
|             await asyncio.sleep(0)
 | |
| 
 | |
|         async def main():
 | |
|             await asyncio.gather(c1(), c2())
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_deep[0], [
 | |
|             'T<anon>',
 | |
|             ['s capture_test_stack', 'a deep', 'a c1'],
 | |
|             [
 | |
|                 ['T<anon>', ['a main', 'a test_stack_gather'], []]
 | |
|             ]
 | |
|         ])
 | |
| 
 | |
|     async def test_stack_shield(self):
 | |
| 
 | |
|         stack_for_shield = None
 | |
| 
 | |
|         async def deep():
 | |
|             await asyncio.sleep(0)
 | |
|             nonlocal stack_for_shield
 | |
|             stack_for_shield = capture_test_stack()
 | |
| 
 | |
|         async def c1():
 | |
|             await asyncio.sleep(0)
 | |
|             await deep()
 | |
| 
 | |
|         async def main():
 | |
|             await asyncio.shield(c1())
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_shield[0], [
 | |
|             'T<anon>',
 | |
|             ['s capture_test_stack', 'a deep', 'a c1'],
 | |
|             [
 | |
|                 ['T<anon>', ['a main', 'a test_stack_shield'], []]
 | |
|             ]
 | |
|         ])
 | |
| 
 | |
|     async def test_stack_timeout(self):
 | |
| 
 | |
|         stack_for_inner = None
 | |
| 
 | |
|         async def inner():
 | |
|             await asyncio.sleep(0)
 | |
|             nonlocal stack_for_inner
 | |
|             stack_for_inner = capture_test_stack()
 | |
| 
 | |
|         async def c1():
 | |
|             async with asyncio.timeout(1):
 | |
|                 await asyncio.sleep(0)
 | |
|                 await inner()
 | |
| 
 | |
|         async def main():
 | |
|             await asyncio.shield(c1())
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_inner[0], [
 | |
|             'T<anon>',
 | |
|             ['s capture_test_stack', 'a inner', 'a c1'],
 | |
|             [
 | |
|                 ['T<anon>', ['a main', 'a test_stack_timeout'], []]
 | |
|             ]
 | |
|         ])
 | |
| 
 | |
|     async def test_stack_wait(self):
 | |
| 
 | |
|         stack_for_inner = None
 | |
| 
 | |
|         async def inner():
 | |
|             await asyncio.sleep(0)
 | |
|             nonlocal stack_for_inner
 | |
|             stack_for_inner = capture_test_stack()
 | |
| 
 | |
|         async def c1():
 | |
|             async with asyncio.timeout(1):
 | |
|                 await asyncio.sleep(0)
 | |
|                 await inner()
 | |
| 
 | |
|         async def c2():
 | |
|             for i in range(3):
 | |
|                 await asyncio.sleep(0)
 | |
| 
 | |
|         async def main(t1, t2):
 | |
|             while True:
 | |
|                 _, pending = await asyncio.wait([t1, t2])
 | |
|                 if not pending:
 | |
|                     break
 | |
| 
 | |
|         t1 = asyncio.create_task(c1())
 | |
|         t2 = asyncio.create_task(c2())
 | |
|         try:
 | |
|             await main(t1, t2)
 | |
|         finally:
 | |
|             await t1
 | |
|             await t2
 | |
| 
 | |
|         self.assertEqual(stack_for_inner[0], [
 | |
|             'T<anon>',
 | |
|             ['s capture_test_stack', 'a inner', 'a c1'],
 | |
|             [
 | |
|                 ['T<anon>',
 | |
|                     ['a _wait', 'a wait', 'a main', 'a test_stack_wait'],
 | |
|                     []
 | |
|                 ]
 | |
|             ]
 | |
|         ])
 | |
| 
 | |
|     async def test_stack_task(self):
 | |
| 
 | |
|         stack_for_inner = None
 | |
| 
 | |
|         async def inner():
 | |
|             await asyncio.sleep(0)
 | |
|             nonlocal stack_for_inner
 | |
|             stack_for_inner = capture_test_stack()
 | |
| 
 | |
|         async def c1():
 | |
|             await inner()
 | |
| 
 | |
|         async def c2():
 | |
|             await asyncio.create_task(c1(), name='there there')
 | |
| 
 | |
|         async def main():
 | |
|             await c2()
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_inner[0], [
 | |
|             'T<there there>',
 | |
|             ['s capture_test_stack', 'a inner', 'a c1'],
 | |
|             [['T<anon>', ['a c2', 'a main', 'a test_stack_task'], []]]
 | |
|         ])
 | |
| 
 | |
|     async def test_stack_future(self):
 | |
| 
 | |
|         stack_for_fut = None
 | |
| 
 | |
|         async def a2(fut):
 | |
|             await fut
 | |
| 
 | |
|         async def a1(fut):
 | |
|             await a2(fut)
 | |
| 
 | |
|         async def b1(fut):
 | |
|             await fut
 | |
| 
 | |
|         async def main():
 | |
|             nonlocal stack_for_fut
 | |
| 
 | |
|             fut = asyncio.Future()
 | |
|             async with asyncio.TaskGroup() as g:
 | |
|                 g.create_task(a1(fut), name="task A")
 | |
|                 g.create_task(b1(fut), name='task B')
 | |
| 
 | |
|                 for _ in range(5):
 | |
|                     # Do a few iterations to ensure that both a1 and b1
 | |
|                     # await on the future
 | |
|                     await asyncio.sleep(0)
 | |
| 
 | |
|                 stack_for_fut = capture_test_stack(fut=fut)
 | |
|                 fut.set_result(None)
 | |
| 
 | |
|         await main()
 | |
| 
 | |
|         self.assertEqual(stack_for_fut[0],
 | |
|             ['F',
 | |
|             [],
 | |
|             [
 | |
|                 ['T<task A>',
 | |
|                     ['a a2', 'a a1'],
 | |
|                     [['T<anon>', ['a test_stack_future'], []]]
 | |
|                 ],
 | |
|                 ['T<task B>',
 | |
|                     ['a b1'],
 | |
|                     [['T<anon>', ['a test_stack_future'], []]]
 | |
|                 ],
 | |
|             ]]
 | |
|         )
 | |
| 
 | |
|         self.assertTrue(stack_for_fut[1].startswith('* Future(id='))
 | |
| 
 | |
| 
 | |
| @unittest.skipIf(
 | |
|     not hasattr(asyncio.futures, "_c_future_add_to_awaited_by"),
 | |
|     "C-accelerated asyncio call graph backend missing",
 | |
| )
 | |
| class TestCallStackC(CallStackTestBase, unittest.IsolatedAsyncioTestCase):
 | |
|     def setUp(self):
 | |
|         futures = asyncio.futures
 | |
|         tasks = asyncio.tasks
 | |
| 
 | |
|         self._Future = asyncio.Future
 | |
|         asyncio.Future = futures.Future = futures._CFuture
 | |
| 
 | |
|         self._Task = asyncio.Task
 | |
|         asyncio.Task = tasks.Task = tasks._CTask
 | |
| 
 | |
|         self._future_add_to_awaited_by = asyncio.future_add_to_awaited_by
 | |
|         futures.future_add_to_awaited_by = futures._c_future_add_to_awaited_by
 | |
|         asyncio.future_add_to_awaited_by = futures.future_add_to_awaited_by
 | |
| 
 | |
|         self._future_discard_from_awaited_by = asyncio.future_discard_from_awaited_by
 | |
|         futures.future_discard_from_awaited_by = futures._c_future_discard_from_awaited_by
 | |
|         asyncio.future_discard_from_awaited_by = futures.future_discard_from_awaited_by
 | |
| 
 | |
|         self._current_task = asyncio.current_task
 | |
|         asyncio.current_task = asyncio.tasks.current_task = tasks._c_current_task
 | |
| 
 | |
|     def tearDown(self):
 | |
|         futures = asyncio.futures
 | |
|         tasks = asyncio.tasks
 | |
| 
 | |
|         futures.future_discard_from_awaited_by = self._future_discard_from_awaited_by
 | |
|         asyncio.future_discard_from_awaited_by = self._future_discard_from_awaited_by
 | |
|         del self._future_discard_from_awaited_by
 | |
| 
 | |
|         futures.future_add_to_awaited_by = self._future_add_to_awaited_by
 | |
|         asyncio.future_add_to_awaited_by = self._future_add_to_awaited_by
 | |
|         del self._future_add_to_awaited_by
 | |
| 
 | |
|         asyncio.Task = self._Task
 | |
|         tasks.Task = self._Task
 | |
|         del self._Task
 | |
| 
 | |
|         asyncio.Future = self._Future
 | |
|         futures.Future = self._Future
 | |
|         del self._Future
 | |
| 
 | |
|         asyncio.current_task = asyncio.tasks.current_task = self._current_task
 | |
| 
 | |
| 
 | |
| @unittest.skipIf(
 | |
|     not hasattr(asyncio.futures, "_py_future_add_to_awaited_by"),
 | |
|     "Pure Python asyncio call graph backend missing",
 | |
| )
 | |
| class TestCallStackPy(CallStackTestBase, unittest.IsolatedAsyncioTestCase):
 | |
|     def setUp(self):
 | |
|         futures = asyncio.futures
 | |
|         tasks = asyncio.tasks
 | |
| 
 | |
|         self._Future = asyncio.Future
 | |
|         asyncio.Future = futures.Future = futures._PyFuture
 | |
| 
 | |
|         self._Task = asyncio.Task
 | |
|         asyncio.Task = tasks.Task = tasks._PyTask
 | |
| 
 | |
|         self._future_add_to_awaited_by = asyncio.future_add_to_awaited_by
 | |
|         futures.future_add_to_awaited_by = futures._py_future_add_to_awaited_by
 | |
|         asyncio.future_add_to_awaited_by = futures.future_add_to_awaited_by
 | |
| 
 | |
|         self._future_discard_from_awaited_by = asyncio.future_discard_from_awaited_by
 | |
|         futures.future_discard_from_awaited_by = futures._py_future_discard_from_awaited_by
 | |
|         asyncio.future_discard_from_awaited_by = futures.future_discard_from_awaited_by
 | |
| 
 | |
|         self._current_task = asyncio.current_task
 | |
|         asyncio.current_task = asyncio.tasks.current_task = tasks._py_current_task
 | |
| 
 | |
| 
 | |
|     def tearDown(self):
 | |
|         futures = asyncio.futures
 | |
|         tasks = asyncio.tasks
 | |
| 
 | |
|         futures.future_discard_from_awaited_by = self._future_discard_from_awaited_by
 | |
|         asyncio.future_discard_from_awaited_by = self._future_discard_from_awaited_by
 | |
|         del self._future_discard_from_awaited_by
 | |
| 
 | |
|         futures.future_add_to_awaited_by = self._future_add_to_awaited_by
 | |
|         asyncio.future_add_to_awaited_by = self._future_add_to_awaited_by
 | |
|         del self._future_add_to_awaited_by
 | |
| 
 | |
|         asyncio.Task = self._Task
 | |
|         tasks.Task = self._Task
 | |
|         del self._Task
 | |
| 
 | |
|         asyncio.Future = self._Future
 | |
|         futures.Future = self._Future
 | |
|         del self._Future
 | |
| 
 | |
|         asyncio.current_task = asyncio.tasks.current_task = self._current_task
 | 
