wizard/shell.py

   1 import subprocess
   2 import sys
   3 import os
   4 import Queue
   5 import threading
   6
   7 import wizard as _wizard
   8 from wizard import util
   9
  10 wizard = sys.argv[0]
  11
  12 class CallError(_wizard.Error):
  13     def __init__(self, code, args, stdout, stderr):
  14         self.code = code
  15         self.args = args
  16         self.stdout = stdout
  17         self.stderr = stderr
  18     def __str__(self):
  19         return "CallError [%d]" % self.code
  20
  21 class PythonCallError(CallError):
  22     def __init__(self, code, args, stdout, stderr):
  23         self.name = util.get_exception_name(stderr)
  24         CallError.__init__(self, code, args, stdout, stderr)
  25     def __str__(self):
  26         return "PythonCallError [%s]" % self.name
  27
  28 def is_python(args):
  29     return args[0] == "python" or args[0] == wizard
  30
  31 class Shell(object):
  32     """An advanced shell, with the ability to do dry-run and log commands"""
  33     def __init__(self, logger = False, dry = False):
  34         """ `logger`    The logger
  35             `dry`       Don't run any commands, just print them"""
  36         self.logger = logger
  37         self.dry = dry
  38     def call(self, *args, **kwargs):
  39         kwargs.setdefault("python", None)
  40         if self.dry or self.logger:
  41             self.logger.info("Running `" + ' '.join(args) + "`")
  42         if self.dry:
  43             return
  44         if kwargs["python"] is None and is_python(args):
  45             kwargs["python"] = True
  46         proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  47         if hasattr(self, "async"):
  48             self.async(proc, args, **kwargs)
  49             return proc
  50         stdout, stderr = proc.communicate()
  51         self.log(stdout, stderr)
  52         if proc.returncode:
  53             if kwargs["python"]: eclass = PythonCallError
  54             else: eclass = CallError
  55             raise eclass(proc.returncode, args, stdout, stderr)
  56         return (stdout, stderr)
  57     def log(self, stdout, stderr):
  58         if self.logger and stdout:
  59             self.logger.debug("STDOUT: " + stdout)
  60         if self.logger and stderr:
  61             self.logger.debug("STDERR: " + stderr)
  62     def callAsUser(self, *args, **kwargs):
  63         user = kwargs.pop("user", None)
  64         kwargs.setdefault("python", is_python(args))
  65         if not user: return self.call(*args, **kwargs)
  66         return self.call("sudo", "-u", user, *args, **kwargs)
  67
  68 class ParallelShell(Shell):
  69     """Commands are queued here, and executed in parallel (with
  70     threading) in accordance with the maximum number of allowed
  71     subprocesses, and result in callback execution when they finish."""
  72     def __init__(self, logger = False, dry = False, max = 10):
  73         super(ParallelShell, self).__init__(logger=logger,dry=dry)
  74         self.running = {}
  75         self.max = max # maximum of commands to run in parallel
  76     def async(self, proc, args, python, on_success, on_error):
  77         """Gets handed a subprocess.Proc object from our deferred
  78         execution"""
  79         self.running[proc.pid] = (proc, args, python, on_success, on_error)
  80     def wait(self):
  81         # bail out immediately on initial ramp up
  82         if len(self.running) < self.max: return
  83         # now, wait for open pids.
  84         try:
  85             pid, status = os.waitpid(-1, 0)
  86         except OSError as e:
  87             if e.errno == errno.ECHILD: return
  88             raise e
  89         # ooh, zombie process. reap it
  90         proc, args, python, on_success, on_error = self.running.pop(pid)
  91         # XXX: this is slightly dangerous; should actually use
  92         # temporary files
  93         stdout = proc.stdout.read()
  94         stderr = proc.stderr.read()
  95         self.log(stdout, stderr)
  96         if status:
  97             if python: eclass = PythonCallError
  98             else: eclass = CallError
  99             on_error(eclass(proc.returncode, args, stdout, stderr))
 100             return
 101         on_success(stdout, stderr)
 102     def join(self):
 103         """Waits for all of our subprocesses to terminate."""
 104         try:
 105             while os.waitpid(-1, 0):
 106                 pass
 107         except OSError as e:
 108             if e.errno == errno.ECHILD: return
 109             raise e
 110
 111 class DummyParallelShell(ParallelShell):
 112     """Same API as ParallelShell, but doesn't actually parallelize (by
 113     using only one thread)"""
 114     def __init__(self, logger = False, dry = False):
 115         super(DummyParallelShell, self).__init__(logger, dry, max=1)
 116