Post

Python Tip #293 (of 365): Write context managers with contextlib.contextmanager. 🧵 Instead of making a class with __enter__ and __exit__ methods, we can make a context manager out of a function: #Python #DailyPythonTip

Trey Hunner

@trey.io

@contextmanager def set_env_var(name, new_value): original_value = os.environ.get(name) os.environ[name] = new_value try: yield finally: if original_value is None: del os.environ[name] else: os.environ[name] = original_value

20/10/26, 3:03 pm0 reposts1 like

Trey Hunner

@trey.io

We can use it in a "with" block, like any other context manager: >>> with set_env_var('USER', 'akin'): ... print(os.environ['USER']) ... akin The code BEFORE the "yield" runs when the "with" block is entered and the code AFTER runs when the block exits.

Trey Hunner

@trey.io

The "try"/"finally" makes sure the environment variable is restored even if an exception is raised in the "with" block. The generator function you decorate must only "yield" one time: that "yield" shouldn't be in a loop and there shouldn't be more than one of them.

Trey Hunner

@trey.io

To customize what your context manager returns to its "with" block, pass a value to the "yield" statement. The contextmanager decorator is great for SIMPLE context managers. If your __enter__ or __exit__ methods need a lot of logic, using a class might be easier to understand.