Put each code in its own dir; save args too

nilmrun/processmanager.py

@@ -184,9 +184,15 @@ class ProcessManager(object):
     """Track and manage a collection of Process objects"""
     def __init__(self):
         self.processes = {}
-        self.tmpfiles = {}
-        self.tmpdir = tempfile.mkdtemp(prefix = "nilmrun-usercode-")
-        atexit.register(shutil.rmtree, self.tmpdir)
+        self.tmpdirs = {}
+
+    def _cleanup_tmpdir(self, pid):
+        if pid in self.tmpdirs:
+            try:
+                shutil.rmtree(self.tmpdirs[pid])
+            except OSError: # pragma: no cover
+                pass
+            del self.tmpdirs[pid]
 
     def __iter__(self):
         return iter(self.processes.keys())
@@ -199,15 +205,33 @@ class ProcessManager(object):
         executed.  The arguments, which must be strings, will be
         accessible in the code as sys.argv[1:]."""
         # The easiest way to do this, by far, is to just write the
-        # code to a file.
-        (fd, path) = tempfile.mkstemp(prefix = "nilmrun-usercode-",
-                                      suffix = ".py", dir=self.tmpdir)
-        with os.fdopen(fd, 'w') as f:
-            f.write(code)
-        argv = [ sys.executable, "-B", "-s", "-u", path ] + args
-        pid = self.run_command(argv)
-        self.tmpfiles[pid] = path
-        return pid
+        # code to a file.  Make a directory to put it in.
+        tmpdir = tempfile.mkdtemp(prefix = "nilmrun-usercode-")
+        try:
+            # Write the code
+            codepath = os.path.join(tmpdir, "usercode.py")
+            with open(codepath, "w") as f:
+                f.write(code)
+            # Save the args too, for debugging purposes
+            with open(os.path.join(tmpdir, "args.txt"), "w") as f:
+                f.write(repr(args))
+
+            # Run the code
+            argv = [ sys.executable, "-B", "-s", "-u", codepath ] + args
+            pid = self.run_command(argv)
+
+            # Save the temp dir
+            self.tmpdirs[pid] = tmpdir
+            tmpdir = None  # Don't need to remove it anymore
+
+            return pid
+        finally:
+            # Clean up tempdir if we didn't finish
+            if tmpdir is not None:
+                try:
+                    shutil.rmtree(tmpdir)
+                except OSError: # pragma: no cover
+                    pass
 
     def run_command(self, argv):
         """Execute a command line program"""
@@ -219,12 +243,7 @@ class ProcessManager(object):
         return self.processes[pid].terminate()
 
     def remove(self, pid):
-        if pid in self.tmpfiles:
-            try:
-                os.unlink(self.tmpfiles[pid])
-            except OSError: # pragma: no cover
-                pass
-            del self.tmpfiles[pid]
+        self._cleanup_tmpdir(pid)
         del self.processes[pid]
 
     def get_info(self):