1 files changed, 62 insertions, 0 deletions
diff --git a/pre_commit/envcontext.py b/pre_commit/envcontext.py
new file mode 100644
index 0000000..1f816ce
--- /dev/null
+++ b/pre_commit/envcontext.py
@@ -0,0 +1,62 @@
+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, 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)