asyncio TaskGroup throwing TypeError when returning lists

Question:

I am trying to use the new TaskGroup object in Python 3.11 to return lists, and I am struggling to find the correct syntax. There isn’t much example code out there yet for this new feature, so I figured I’d ask.

import asyncio

async def squares(nums):
    return await [n**2 for n in nums]

async def squares_tasks(nums1, nums2):
    async with asyncio.TaskGroup as tg:
        squares1 = tg.create_task(squares(nums1))
        squares2 = tg.create_task(squares(nums2))
        return squares1, squares2

nums1, nums2 = [1, 2, 3], [4, 5, 6]
squares1, squares2 = asyncio.run(squares_tasks(nums1, nums2))

This throws the following error:
TypeError: 'type' object does not support the context manager protocol

I would appreciate any pointers for what I am doing wrong, thanks.

Asked By: sha2fiddy

||

Answers:

There are multiple errors in your program. Here is a correct one:

import asyncio


async def squares(nums):
    return [n**2 for n in nums]


async def squares_tasks(nums1, nums2):
    async with asyncio.TaskGroup() as tg:
        squares1 = tg.create_task(squares(nums1))
        squares2 = tg.create_task(squares(nums2))

    r1 = squares1.result()
    r2 = squares2.result()

    return r1, r2


if __name__ == "__main__":
    nums1, nums2 = [1, 2, 3], [4, 5, 6]
    squares1, squares2 = asyncio.run(squares_tasks(nums1, nums2))
    print(squares1, squares2)

First, you must wait for the tasks to finish, that is, you must use squares1 and squares2 outside of the async-with statement. Second, you must use the .result() method to get the actual result of the task.

Please, also note that there is nothing that is async in your squares function so you will not get any benefits from using AsyncIO here.

Answered By: Louis Lac