bpo-46994: Accept explicit contextvars.Context in asyncio create_task() API (GH-31837)

This commit is contained in:
Andrew Svetlov 2022-03-14 13:54:13 +02:00 committed by GitHub
parent 2153daf0a0
commit 9523c0d84f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 209 additions and 65 deletions

View file

@ -95,8 +95,8 @@ class BaseTaskTests:
Task = None
Future = None
def new_task(self, loop, coro, name='TestTask'):
return self.__class__.Task(coro, loop=loop, name=name)
def new_task(self, loop, coro, name='TestTask', context=None):
return self.__class__.Task(coro, loop=loop, name=name, context=context)
def new_future(self, loop):
return self.__class__.Future(loop=loop)
@ -2527,6 +2527,90 @@ class BaseTaskTests:
self.assertEqual(cvar.get(), -1)
def test_context_4(self):
cvar = contextvars.ContextVar('cvar')
async def coro(val):
await asyncio.sleep(0)
cvar.set(val)
async def main():
ret = []
ctx = contextvars.copy_context()
ret.append(ctx.get(cvar))
t1 = self.new_task(loop, coro(1), context=ctx)
await t1
ret.append(ctx.get(cvar))
t2 = self.new_task(loop, coro(2), context=ctx)
await t2
ret.append(ctx.get(cvar))
return ret
loop = asyncio.new_event_loop()
try:
task = self.new_task(loop, main())
ret = loop.run_until_complete(task)
finally:
loop.close()
self.assertEqual([None, 1, 2], ret)
def test_context_5(self):
cvar = contextvars.ContextVar('cvar')
async def coro(val):
await asyncio.sleep(0)
cvar.set(val)
async def main():
ret = []
ctx = contextvars.copy_context()
ret.append(ctx.get(cvar))
t1 = asyncio.create_task(coro(1), context=ctx)
await t1
ret.append(ctx.get(cvar))
t2 = asyncio.create_task(coro(2), context=ctx)
await t2
ret.append(ctx.get(cvar))
return ret
loop = asyncio.new_event_loop()
try:
task = self.new_task(loop, main())
ret = loop.run_until_complete(task)
finally:
loop.close()
self.assertEqual([None, 1, 2], ret)
def test_context_6(self):
cvar = contextvars.ContextVar('cvar')
async def coro(val):
await asyncio.sleep(0)
cvar.set(val)
async def main():
ret = []
ctx = contextvars.copy_context()
ret.append(ctx.get(cvar))
t1 = loop.create_task(coro(1), context=ctx)
await t1
ret.append(ctx.get(cvar))
t2 = loop.create_task(coro(2), context=ctx)
await t2
ret.append(ctx.get(cvar))
return ret
loop = asyncio.new_event_loop()
try:
task = loop.create_task(main())
ret = loop.run_until_complete(task)
finally:
loop.close()
self.assertEqual([None, 1, 2], ret)
def test_get_coro(self):
loop = asyncio.new_event_loop()
coro = coroutine_function()