Coverage for .nox/test-3-9/lib/python3.9/site-packages/nskit/common/contextmanagers/env.py: 91%

1"""Context manager for running in a specific directory.""" 

2from contextlib import ContextDecorator 

3import os 

4from typing import Dict, List, Optional 

5 

6from nskit._logging import logger_factory 

7 

8 

9class Env(ContextDecorator): 

10 """Context manager for managing environment variables. 

11 

12 The optional arguments can provide either an exhaustive set of environment values, values to override or values to remove. 

13 """ 

14 

15 def __init__( 

16 self, 

17 environ: Optional[Dict[str, str]] = None, 

18 override: Optional[Dict[str, str]] = None, 

19 remove: Optional[List[str]] = None 

20 ): 

21 """Initialise the context manager. 

22 

23 The parameters are applied in the following order (so can be combined): 1st - environ, 2nd - override, 3rd - remove 

24 

25 Keyword Args: 

26 environ (Optional[Dict[str, str]]): an exhaustive set of environment values to set (replaces overall os.environ contents) 

27 override (Optional[Dict[str, str]]): a set of environment values to either override or set (replaces values in existing os.environ) 

28 remove (Optional[List[str]]): a set of environment values to remove (removes values if found in os.environ ) 

29 """ 

30 if environ is not None and not isinstance(environ, dict): 

31 raise TypeError('environ should be a dict') 

32 if override is not None and not isinstance(override, dict): 

33 raise TypeError('override should be a dict') 

34 if remove is not None and not isinstance(remove, (list, tuple, set)): 

35 raise TypeError('remove should be a (list, tuple, set)') 

36 self._environ = environ 

37 self._override = override 

38 self._remove = remove 

39 self._original = None 

40 

41 def __enter__(self): 

42 """Change to the target environment variables.""" 

43 # Handling circular imports with LoggingConfig 

44 logger = logger_factory.get_logger(__name__) 

45 if self._environ or self._override or self._remove: 

46 self._original = os.environ.copy() 

47 if self._environ is not None: 

48 # Here we pop all keys from it and then update 

49 os.environ.clear() 

50 os.environ.update(self._environ) 

51 if self._override: 

52 os.environ.update(self._override) 

53 if self._remove: 

54 for key in list(self._remove): 

55 os.environ.pop(key, None) 

56 logger.info('Changing env variables') 

57 logger.debug(f'New env variables: {os.environ}') 

58 else: 

59 logger.info('No arguments set (environ, override, remove)') 

60 

61 def __exit__(self, *args, **kwargs): # noqa: U100 

62 """Reset to the original environment variables.""" 

63 if self._original: 

64 os.environ.clear() 

65 os.environ.update(self._original.copy())