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.

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

Your post must include a link to this post.

How does this work?
Write a response on your own website and include a link to this post. Then submit your response’s URL here. Your response may appear here after moderation. To update or remove it, edit or delete your original post and submit the same URL again. Learn more about Webmentions.