62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import enum
|
|
import os
|
|
from collections.abc import Generator
|
|
from collections.abc import MutableMapping
|
|
from typing import NamedTuple
|
|
from typing import Union
|
|
|
|
_Unset = enum.Enum('_Unset', 'UNSET')
|
|
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: MutableMapping[str, str] | None = None,
|
|
) -> Generator[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)
|