|
|
@@ -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): |
|
|
|