1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
|
import contextlib
import enum
import os
from typing import Generator
from typing import MutableMapping
from typing import NamedTuple
from typing import Optional
from typing import Tuple
from typing import Union
class _Unset(enum.Enum):
UNSET = 1
UNSET = _Unset.UNSET
class Var(NamedTuple):
name: str
default: str = ''
SubstitutionT = Tuple[Union[str, Var], ...]
ValueT = Union[str, _Unset, SubstitutionT]
PatchesT = Tuple[Tuple[str, ValueT], ...]
def format_env(parts: SubstitutionT, env: MutableMapping[str, str]) -> str:
return ''.join(
env.get(part.name, part.default) if isinstance(part, Var) else part
for part in parts
)
@contextlib.contextmanager
def envcontext(
patch: PatchesT,
_env: Optional[MutableMapping[str, str]] = None,
) -> Generator[None, None, None]:
"""In this context, `os.environ` is modified according to `patch`.
`patch` is an iterable of 2-tuples (key, value):
`key`: string
`value`:
- string: `environ[key] == value` inside the context.
- UNSET: `key not in environ` inside the context.
- template: A template is a tuple of strings and Var which will be
replaced with the previous environment
"""
env = os.environ if _env is None else _env
before = dict(env)
for k, v in patch:
if v is UNSET:
env.pop(k, None)
elif isinstance(v, tuple):
env[k] = format_env(v, before)
else:
env[k] = v
try:
yield
finally:
env.clear()
env.update(before)
|
