Wrap Context Manager (as nested) in a Python function

To make my code more elegant, I need to wrap a context manager with initialization code as a function in Python. This is definitely possible, but it took me some time to find the most elegant way to do this.

Generally speaking, you will want to enter all contexts when using the decorator @contextlib.contextmanager or @contextlib.asynccontextmanager. When the end user uses with my_function() as a:, everything inside the with block has been inside the nested contexts. When Python gets out of the end user's with block, it should also run all related __exit__s in the wrapper function. See example (requires Python 3.7+ probably and writes test.txt; you can test it on repl.it):

import asyncio
import contextlib
import aiofiles
import typing

@contextlib.asynccontextmanager
async def my_context_manager() -> typing.ContextManager[aiofiles.threadpool.AsyncFileIO]:
    # init
    filename = 'test.txt'

    async with aiofiles.open(filename, 'w+') as file:
        # you can still override or interact with `file` if needed
        yield file

async def main():
    # end user
    myfile: aiofiles.threadpool.AsyncFileIO
    async with my_context_manager() as myfile:
        print(await myfile.write("12\n"))
        await myfile.seek(0)
        print(await myfile.readline())

asyncio.run(main())

When I rewrite the code above I feel it so easy and natural. It really took me some time to realize how this works.

发表评论

电子邮件地址不会被公开。 必填项已用*标注