mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-04-23 03:00:31 -04:00
26 lines
807 B
Python
26 lines
807 B
Python
import os
|
|
from contextlib import contextmanager
|
|
|
|
|
|
@contextmanager
|
|
def unset_env_var(env_var: str):
|
|
"""Context manager that unsets an environment variable, re-setting it to the original value when the context is exited."""
|
|
prev_value = os.environ.get(env_var, None)
|
|
if prev_value is not None:
|
|
del os.environ[env_var]
|
|
yield
|
|
if prev_value is not None:
|
|
os.environ[env_var] = prev_value
|
|
|
|
|
|
@contextmanager
|
|
def set_env_var(env_var: str, value: str):
|
|
"""Context manager that sets an environment variable, re-setting it to the original value when the context is exited."""
|
|
prev_value = os.environ.get(env_var, None)
|
|
os.environ[env_var] = value
|
|
yield
|
|
if prev_value is not None:
|
|
os.environ[env_var] = prev_value
|
|
else:
|
|
del os.environ[env_var]
|