util: os: prepend_to_path

Signed-off-by: John Andersen <[email protected]>

Author: pdxjohnny

dffml/util/os.py

@@ -1,8 +1,8 @@
 import os
-from contextlib import contextmanager
+import contextlib
 
 
-@contextmanager
+@contextlib.contextmanager
 def chdir(new_path):
     """
     Context manager to change directroy
@@ -13,3 +13,17 @@ def chdir(new_path):
         yield
     finally:
         os.chdir(old_path)
+
+
+@contextlib.contextmanager
+def prepend_to_path(*args: str):
+    """
+    Prepend all given directories to the ``PATH`` environment variable.
+    """
+    old_path = os.environ.get("PATH", "")
+    # TODO Will this work on Windows?
+    os.environ["PATH"] = ":".join(list(map(str, args)) + old_path.split(":"))
+    try:
+        yield
+    finally:
+        os.environ["PATH"] = old_path

tests/util/test_os.py

@@ -0,0 +1,13 @@
+import os
+import unittest
+
+from dffml.util.os import prepend_to_path
+
+
+class TestOS(unittest.TestCase):
+    def test_prepend_to_path(self):
+        old_path = os.environ["PATH"]
+        with prepend_to_path("jellybeans", "fritos"):
+            self.assertEqual(
+                os.environ["PATH"], "jellybeans:fritos:" + old_path
+            )