Temporary modifying env variables with context manager One function of a program I'm working with benefits from using all available openMP threads. The rest of the program has issues with the env variable OMP_NUM_THREADS being unequal to 1.

I'm currently using the following solution to that, and was wondering whether there were any caveats or room for improvement here. I haven't worked much with context managers yet, and greatly appreciate any help/pointers. Thanks!

import os
import multiprocessing as mp

class MaxOMP:
    def __enter__(self):
        os.environ["OMP_NUM_THREADS"] = str(mp.cpu_count())

    def __exit__(self, exc_type, exc_val, exc_tb):
        os.environ["OMP_NUM_THREADS"] = "1"


def main():
    with MaxOMP():
        run_omp_intensive_job()
Instead of MaxOMP as a name, I like MaxOmp - in general still CamelCase your names even if they include an acronym.

Additionally, there isn't any reason why this needs to be restricted to any particular environment variable. As such, I would probably make this:

class EnvironmentVariableContextManager:

    _base_value = None
    _variable = None

    def __init__(self, variable_name, value):
        self.variable = variable_name
        self.target_value = value

    @property
    def variable(self):
        return self._variable

    @variable.setter
    def variable(self, name):
        self._base_value = os.environ.get(name, None)

        if self._base_value is None:
            raise Exception(f"{name} is not a valid environment variable")

        self._variable = name

    @property
    def original_value(self):
        return self._base_value

    def __enter__(self):
        os.environ[self.variable] = self.target_value

    def __exit__(self, *):
        os.environ[self.variable] = self.original_value


This has a few benefits:


We use properties to properly encapsulate access and add some extra validation (if you don't care about invalid environment variables feel free to remove that)
We restore it to the original value, not some hardcoded value. You could also provide a way to specify an "after" value if that was desirable.


As suggested in the comments, we may have a valid use-case for an environment variable that doesn't currently have a value. In that case, we would remove the exception, and instead do this:

@variable.setter
def variable(self, name):
    self._base_value = os.environ.get(name, None)
    self._variable = name

def __exit__(self, *):
    if self.original_value is None:
        os.environ.pop(self.variable)
    else:
        os.environ[self.variable] = self.original_value


Once we have that, then we just inherit.

class MaxOmpThreads(EnvironmentVariableContextManager):

    def __init__(self, max_threads=None):
        max_threads = max_threads or mp.cpu_count()
        super().__init__("OMP_NUM_THREADS", max_threads)


Another thing to note, is that os.environ only reflects the environment variables at the time of Python startup. If this is not what you want then you'll likely need a different solution.

This also may not be thread-safe; if you might have multiple processes working on something at once, and some need one setting and some need another, then you're likely to run into issues. If your program is overall serial, and the pieces that use OMP are well defined and contained, you're probably fine.