Adding env context handler

pie.py

@@ -4,7 +4,7 @@
 Enables a user to execute predefined tasks that may accept parameters and options from the command line without any other required packages.
 Great for bootstrapping a development environment, and then interacting with it.
 """
-__VERSION__='0.3.0a'
+__VERSION__='0.3.0b'
 
 
 import inspect
@@ -16,7 +16,7 @@
 import types
 
 
-__all__=['task','Parameter','OptionsParameter','options','cmd','cd','pip','venv']
+__all__=['task','Parameter','OptionsParameter','options','cmd','cd','env','pip','venv']
 
 
 # environment constants
@@ -306,6 +306,27 @@ def exit_hook(self):
         os.chdir(self.old_path)
 
 
+class env(CmdContext):
+    """A context class used to execute commands with a different environment"""
+    def __init__(self,env_dict):
+        self.env_dict=env_dict
+
+    def enter_hook(self):
+        self.old_env_dict={k:os.environ.get(k,None) for k in self.env_dict.keys()}
+        self._set_environ(self.env_dict)
+
+    def exit_hook(self):
+        self._set_environ(self.old_env_dict)
+
+    @classmethod
+    def _set_environ(cls,d):
+        for k,v in d.items():
+            if v is None:
+                if k in os.environ:
+                    del os.environ[k]
+            else:
+                os.environ[k]=v
+
 
 # ----------------------------------------
 # Command line functionality

tests/data/env_context/pie_tasks.py

@@ -0,0 +1,17 @@
+import os
+
+from pie import *
+
+
+@task
+def envContext():
+    e={'MY_VAL':'test','HOME':'something','COMPUTERNAME':None}
+
+    def _print_env():
+        for k in e:
+            print('{}: {}'.format(k,os.environ.get(k,'unset')))
+
+    _print_env()
+    with env(e):
+        _print_env()
+    _print_env()

tests/test_context_manager.py

@@ -21,3 +21,19 @@ def test_cd_context(pie,capsys,pie_tasks_path,pie_mock_cmd):
     assert len(pie_mock_cmd.cmds)==0
     out,err=capsys.readouterr()
     assert out=='{}\n{}\n'.format(os.path.dirname(cwd),cwd)
+
+
+@pytest.mark.parametrize('pie_tasks_path',['env_context'],indirect=['pie_tasks_path'])
+def test_env_context(pie,capsys,pie_tasks_path,pie_mock_cmd):
+    pie.main(['envContext'])
+
+    out_vars={'MY_VAL':['test'],'HOME':['something'],'COMPUTERNAME':['unset']}
+    out,err=capsys.readouterr()
+    for l in out.split('\n'):
+        if l=='': continue
+        k,v=l.split(': ',1)
+        out_vars[k].append(v)
+    for k,v in out_vars.items():
+        assert v[0]==v[2]
+        assert v[1]==v[3]
+        assert v[0]!=v[1] # this may fail on *nix for COMPUTERNAME