# 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):

def __exit__(self, exc_type, exc_val, exc_tb):

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):


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.
You don't need multiprocessing as os already have what you need : os.cpu_count