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.