nested "async with" using aiohttp

Question

I would like to create a scheduler class that uses aiohttp to make API calls. I tried this:

import asyncio
import aiohttp

class MySession:
    def __init__(self):
        self.session = None

    async def __aenter__(self):
        async with aiohttp.ClientSession() as session:
            self.session = session
            return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()

async def method1():
    async with MySession() as s:
        async with s.session.get("https://www.google.com") as resp:
            if resp.status == 200:
                print("successful call!")

loop = asyncio.get_event_loop()
loop.run_until_complete(method1())
loop.close()

but this just results in an error: RuntimeError: Session is closed.

A second approach for the __aenter__ function:

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self

works well. Is this a good construct? It doesn't adhere to examples of how to use aiohttp. Also wondering why the first approach isn't working?

Answer 1

You can't use with inside a function and have the context manager remain open, no. The with with aiohttp.ClientSession() as session: block exits as soon as you use return to exit the __aenter__ coroutine!

For the specific case, entering a aiohttp.ClientSession() context manager does nothing but return self. So for that type, just creating the instance and storing it in self.session, and awaiting on self.session.close() suffices here, yes.

The general pattern for a nested asynchronous context manager is to await the __aenter__ and __aexit__ methods of a nested async context manager from your own such methods (and perhaps pass along the exception information):

class MySession:
    def __init__(self):
        self.session = None

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        await self.session.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            return await self.session.__aexit__(exc_type, exc_val, exc_tb)

Technically speaking, you should first assure that there is an actual __aexit__ attribute before entering a nested context manager:

class MySession:
    def __init__(self):
        self.session = None
        self._session_aexit = None

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        self._session_aexit = type(self.session).__aexit__
        await self.session.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            return await self._session_aexit.__aexit__(
                self.session, exc_type, exc_val, exc_tb)

See the official PEP that added the concept.

Answer 2

You can manage that dependency externally:

import asyncio
import aiohttp

class MySession:

    def __init__(self, client_session):
        self.session = client_session

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass

async def method1():
    async with aiohttp.ClientSession() as client_session:
        async with MySession(client_session) as s:
            async with s.session.get("https://www.google.com") as resp:
                if resp.status == 200:
                    print("successful call!")

asyncio.run(method1())

When that async with chain becomes too ridiculous you can use AsyncExitStack:

from contextlib import AsyncExitStack

async def method1():
    async with AsyncExitStack() as stack:
        client_session = await stack.enter_async_context(aiohttp.ClientSession())
        s = await stack.enter_async_context(MySession(client_session))
        async with s.session.get("https://www.google.com") as resp:
            if resp.status == 200:
                print("successful call!")