ansible: split environment editing into a separate class.

ansible_mitogen/helpers.py

@@ -112,6 +112,17 @@ def module_fixups(mod):
         mod.YumRepo.repofile = mod.configparser.RawConfigParser()
 
 
+class TemporaryEnvironment(object):
+    def __init__(self, env=None):
+        self.original = os.environ.copy()
+        self.env = env or {}
+        os.environ.update((k, str(v)) for k, v in self.env.iteritems())
+
+    def revert(self):
+        os.environ.clear()
+        os.environ.update(self.original)
+
+
 def run_module(module, raw_params=None, args=None, env=None):
     """
     Set up the process environment in preparation for running an Ansible
@@ -130,10 +141,8 @@ def run_module(module, raw_params=None, args=None, env=None):
         'ANSIBLE_MODULE_ARGS': args
     })
 
-    if env:
-        original_env = os.environ.copy()
-        os.environ.update((k, str(v)) for k, v in env.iteritems())
-
+    temp_env = TemporaryEnvironment(env)
+    try:
         try:
             mod = __import__(module, {}, {}, [''])
             module_fixups(mod)
@@ -145,10 +154,8 @@ def run_module(module, raw_params=None, args=None, env=None):
             mod.main()
         except (Exit, ModuleError), e:
             result = json.dumps(e.dct)
-
-    if env:
-        os.environ.clear()
-        os.environ.update(original_env)
+    finally:
+        temp_env.revert()
 
     return result