Add envcontext helper

pre_commit/envcontext.py

@@ -0,0 +1,54 @@
+from __future__ import absolute_import
+from __future__ import unicode_literals
+
+import collections
+import contextlib
+import os
+
+from pre_commit import five
+
+
+UNSET = collections.namedtuple('UNSET', ())()
+
+
+Var = collections.namedtuple('Var', ('name', 'default'))
+setattr(Var.__new__, five.defaults_attr, ('',))
+
+
+def format_env(parts, env):
+    return ''.join(
+        env.get(part.name, part.default)
+        if isinstance(part, Var)
+        else part
+        for part in parts
+    )
+
+
+@contextlib.contextmanager
+def envcontext(patch, _env=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 = env.copy()
+
+    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)

pre_commit/five.py

@@ -12,6 +12,8 @@ if PY2:  # pragma: no cover (PY2 only)
             return s
         else:
             return s.encode('UTF-8')
+
+    defaults_attr = 'func_defaults'
 else:  # pragma: no cover (PY3 only)
     text = str
 
@@ -21,6 +23,8 @@ else:  # pragma: no cover (PY3 only)
         else:
             return s.decode('UTF-8')
 
+    defaults_attr = '__defaults__'
+
 
 def to_text(s):
     return s if isinstance(s, text) else s.decode('UTF-8')