Browse Source

Put each code in its own dir; save args too

tags/nilmrun-1.2.5
Jim Paris 10 years ago
parent
commit
bed26e099b
1 changed files with 37 additions and 18 deletions
  1. +37
    -18
      nilmrun/processmanager.py

+ 37
- 18
nilmrun/processmanager.py View File

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


Loading…
Cancel
Save