Coverage for src/mlopus/utils/env_utils.py: 100%

1import contextlib 

2import os 

3import threading 

4from typing import Dict, ContextManager 

5 

6_env_lock = threading.Lock() 

7 

8 

9@contextlib.contextmanager 

10def using_env_vars(_vars: Dict[str, str]) -> ContextManager[None]: 

11 """Thread-safely set env vars inside context and reset them on exit.""" 

12 with _env_lock: 

13 original_env = dict(os.environ) 

14 try: 

15 os.environ.update(_vars) 

16 yield None 

17 finally: 

18 os.environ = original_env