Skip to content

Contextlib

  • provides utilities for common tasks involving the with statement.

What it provides?

  • contextmanager
  • asynccontextmanager
  • closing
  • aclosing
  • nullcontext
  • suppress
  • redirect_stdout
  • redirect_stderr
  • chdir
  • ContextDecorator

Implementing a simple Context Manager (using class)

  • A class needs to implement __enter__ and __exit__ methods to be used as a context manager.

  • __enter__ method is called when the with block is entered. It doesn't take any arguments.

  • __exit__ method is called when the with block is exited. It takes three arguments:
    • exc_type: exception type
    • exc_value: exception value
    • traceback: traceback object
  • If no exception occurs, all three arguments are None. Else, they contain the exception details.
class MyCtx:
    def __init__(self, num):
        self.num = num

    def __enter__(self):
        print("enter block executing")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("exit block executing")


if __name__ == "__main__":
    with MyCtx(4) as c:
        print(f"value of {c.num=}")

    print("-"*40)

contextmanager decorator

  • A simpler way to create a context manager is to use the contextmanager decorator from the contextlib module.
  • You don't need to implement a whole new class just to be able to use with statements with it.
  • function decorated with @contextmanager should yield the resource that needs to be managed.
  • Use try, except and finally blocks.
  • exception raised inside the with block can be caught in the except block.
  • finally block is used to clean up the resource, after the with block is exited.
from contextlib import contextmanager

@contextmanager
def my_ctx(*args, **kwargs):
    try:
        print('Entering context')
        yield 'hello'
    except Exception as e:
        print(f'Caught exception: {e}')
    finally:
        print('Exiting context')

if __name__ == '__main__':
    with my_ctx() as val:
        print(f'Value: {val}')
    print("-"*50)

AsyncContextManager

  • similar to contextmanager, but for async functions.
  • Used with async with statements.
  • To implement, use __aenter__ and __aexit__ methods.
from contextlib import asynccontextmanager

@asynccontextmanager
async def get_connection():
    conn = await acquire_db_connection()
    try:
        yield conn
    finally:
        await release_db_connection(conn)

async def get_all_users():
    async with get_connection() as conn:
        return conn.query('SELECT ...')

Closing

  • closing is a helper function that simplifies the use of resources that need to be closed properly, such as files, sockets, or database connections.
  • It ensures that the resource gets closed automatically when you're done with it, even if an error occurs.

Why Use closing?

  • Some objects don't support the with statement directly (they don't have a __enter__ or __exit__ method).
  • The closing function makes it possible to use these objects in a with block, ensuring proper cleanup.
  • When the block ends, closing calls the resource's close() method automatically
from contextlib import contextmanager, closing

class MyCtx:
    def __init__(self, name):
        self.name = name
        print("myctx function called")
    def close(self):
        print("myctx close function called")


if __name__ == "__main__":
    with closing(MyCtx("deependu")) as m:
        print(m.name)

expected Output:

myctx function called
deependu
myctx close function called

Aclosing

  • aclosing is the async version of closing.

NullContext

  • do-nothing placeholder that can be used with with statement.
  • Used when you don't need to do anything in the __enter__ and __exit__ methods.
  • Mostly used for maintaining consistency in the code.
  • Whatever a function returns, it should return support with statement.
def myfunction(arg, ignore_exceptions=False):
    if ignore_exceptions:
        # Use suppress to ignore all exceptions.
        cm = contextlib.suppress(Exception)
    else:
        # Do not ignore any exceptions, cm has no effect.
        cm = contextlib.nullcontext()
    with cm:
        # Do something
  • in the above code, we wanted to use with block, so we used nullcontext to maintain consistency.

Suppress

  • Return a context manager that suppresses any of the specified exceptions if they occur in the body of a with statement and then resumes execution with the first statement following the end of the with statement.
from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove('somefile.tmp')

## equivalent to
# try:
#     os.remove('somefile.tmp')
# except FileNotFoundError:
#     pass
  • To pass multiple exceptions, use a tuple.
from contextlib import suppress

if __name__ == "__main__":
    with suppress((ValueError, FileNotFoundError)):
        raise ValueError()
        raise FileNotFoundError()

Redirect_stdout

  • Instead of printing to the console, you can redirect the output to a file or a string.
  • To do this, we can modify the sys.stdout object to point to a different file or string.
import sys

sys.stdout = open('output.txt', 'w')
print('This is redirected to a file')
  • But, this is not a good practice, as it can lead to issues with other parts of the code that rely on sys.stdout, if we forget to reset it back to the original value.
  • redirect_stdout provides a context manager that temporarily redirects sys.stdout to a different file or stream.
from contextlib import redirect_stdout

with open('output.txt', 'w') as f:
    with redirect_stdout(f):
        print('This is redirected to a file')

print('This is printed to the console')

Redirect_stderr

  • Similar to redirect_stdout, but for sys.stderr.

chdir

  • This is a simple wrapper around chdir(), it changes the current working directory upon entering and restores the old one on exit.

ContextDecorator

  • Used to create custom decorators that can be used as context managers.
  • @contextmanager is implemented using the same method.
  • Class inheriting from ContextDecorator have to implement __enter__ and __exit__ as normal.
from contextlib import ContextDecorator

class mycontext(ContextDecorator):
    def __enter__(self):
        print('Starting')
        return self

    def __exit__(self, *exc):
        print('Finishing')
        return False

# =============
@mycontext()
def function():
    print('The bit in the middle')

function()
# output:   Starting
#           The bit in the middle
#           Finishing

with mycontext():
    print('The bit in the middle')

# output:   Starting
#           The bit in the middle
#           Finishing

How are they implemented?

nullcontext

class MyNullcontext:
    def __enter__(self):
        return None

    def __exit__(self, exc_type, exc_value, traceback):
        pass

closing

class MyClosing:
    def __enter__(self, some_resource):
        return some_resource

    def __exit__(self, exc_type, exc_value, traceback):
        some_resource.close() # just calls the `close` method of the resource