mirror of
				https://github.com/python/cpython.git
				synced 2025-11-04 11:49:12 +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
 |