Custom context manager for environment variables.

Let's now consider a use-case where you might want to write your own context manager: it's when you're writing unit tests that deal with environment variables.

import os

os.environ['APP_KEY'] = 'demo-app-key'
os.environ['APP_KEY']

Environment variables are typically used as a way to configure settings from outside of a program. You could imagine an environment variable that determines the location of a database that a webapp might use. If you're writing that webapp that's typically something that you would want to unit-test though. And you'd want to make sure that setting an environment variable in one test doens't have any effect on any other tests.

Context Manager for Environment Variables

So let's write a context manager to help us with that.

from contextlib import contextmanager

@contextmanager
def set_env_vars(**kwargs):
    # Remember old values
    previous_values = {k: os.environ[k] for k in kwargs if os.environ.get(k, None)}
    # Set new values
    for k, v in kwargs.items():
        os.environ[k] = v
    yield
    # Put everything back into place
    for k, v in previous_values.items():
        os.environ[k] = v

This context manager will set environment variables based on the **kwargs that are passed and will make sure that these are unset when the manager exits. That also means that a codeblock like this ...

with set_env_vars(APP_KEY="something-else"):
    print(os.environ["APP_KEY"])
print(os.environ["APP_KEY"])

... will output this:

something-else
demo-app-key

This is a great example of a small context manager that might make it easier to deal with your own codebase. While it's always good to check the main libraries for the managers that they provide, it's also a good idea to take a step back and consider what utilities you might be able to write for yourself.