From a5e4a6f346213fc80b480a6768ce60a5b127ecb9 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 1 Apr 2018 21:40:06 +0100
Subject: [PATCH] issue #106: move helpers.get_bytecode() into NewStyleRunner

---
 ansible_mitogen/helpers.py | 10 ----------
 ansible_mitogen/runner.py  | 31 ++++++++++++++++++++-----------
 2 files changed, 20 insertions(+), 21 deletions(-)

diff --git a/ansible_mitogen/helpers.py b/ansible_mitogen/helpers.py
index b107d000..270f3244 100644
--- a/ansible_mitogen/helpers.py
+++ b/ansible_mitogen/helpers.py
@@ -51,9 +51,6 @@ LOG = logging.getLogger(__name__)
 #: Caching of fetched file data.
 _file_cache = {}
 
-#: Caching of compiled new-style module bytecode.
-_bytecode_cache = {}
-
 #: Mapping of job_id<->result dict
 _result_by_job_id = {}
 
@@ -87,13 +84,6 @@ def get_file(context, path):
     return _file_cache[path]
 
 
-def get_bytecode(context, path):
-    if path not in _bytecode_cache:
-        source = get_file(context, path)
-        _bytecode_cache[path] = compile(source, path, 'exec')
-    return _bytecode_cache[path]
-
-
 def run_module(kwargs):
     """
     Set up the process environment in preparation for running an Ansible
diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py
index a6b77977..f5cb1283 100644
--- a/ansible_mitogen/runner.py
+++ b/ansible_mitogen/runner.py
@@ -306,32 +306,41 @@ class NewStyleRunner(ScriptRunner):
     Execute a new-style Ansible module, where Module Replacer-related tricks
     aren't required.
     """
+    #: path => new-style module bytecode.
+    _code_by_path = {}
+
     def setup(self):
         super(NewStyleRunner, self).setup()
         self._stdio = NewStyleStdio(self.args)
         self._argv = TemporaryArgv([self.path])
 
     def revert(self):
-        super(NewStyleRunner, self).revert()
+        self._argv.revert()
         self._stdio.revert()
+        super(NewStyleRunner, self).revert()
 
-    def _get_bytecode(self):
-        """
-        Fetch the module binary from the master if necessary.
-        """
-        return ansible_mitogen.helpers.get_bytecode(
-            context=self.service_context,
-            path=self.path,
-        )
+    def _get_code(self):
+        try:
+            return self._code_by_path[self.path]
+        except KeyError:
+            return self._code_by_path.setdefault(self.path, compile(
+                source=ansible_mitogen.helpers.get_file(
+                    context=self.service_context,
+                    path=self.path,
+                ),
+                filename=self.path,
+                mode='exec',
+                dont_inherit=True,
+            ))
 
     def _run(self):
-        bytecode = self._get_bytecode()
+        code = self._get_code()
         mod = types.ModuleType('__main__')
         d = vars(mod)
         e = None
 
         try:
-            exec bytecode in d, d
+            exec code in d, d
         except SystemExit, e:
             pass