]> scripts.mit.edu Git - wizard.git/blobdiff - wizard/shell.py
Refactor interactive shell and add addenv support.
[wizard.git] / wizard / shell.py
index 9f2d43c7eecdb72b3a726decb6dae02e42055f02..523dea23af6e1e7d1887b2777aceae6e8c1ca83c 100644 (file)
@@ -55,6 +55,7 @@ class Shell(object):
     """
     def __init__(self, dry = False):
         self.dry = dry
+        self.cwd = None
     def call(self, *args, **kwargs):
         """
         Performs a system call.  The actual executable and options should
@@ -75,6 +76,10 @@ class Shell(object):
             specify this if you are using another wrapper around this function).
         :param log: if True, we log the call as INFO, if False, we log the call
             as DEBUG, otherwise, we detect based on ``strip``.
+        :param addenv: mapping of environment variables *to add*
+        :param stdout:
+        :param stderr:
+        :param stdin: a file-type object that will be written to or read from as a pipe.
         :returns: a tuple of strings ``(stdout, stderr)``, or a string ``stdout``
             if ``strip`` is specified.
 
@@ -84,12 +89,15 @@ class Shell(object):
         >>> sh.call("cat", input='Foobar')
         ('Foobar', '')
         """
-        if hasattr(self, "_wait"):
-            self._wait()
+        self._wait()
         kwargs.setdefault("interactive", False)
         kwargs.setdefault("strip", False)
         kwargs.setdefault("python", None)
         kwargs.setdefault("log", None)
+        kwargs.setdefault("stdout", subprocess.PIPE)
+        kwargs.setdefault("stdin", subprocess.PIPE)
+        kwargs.setdefault("stderr", subprocess.PIPE)
+        kwargs.setdefault("addenv", None)
         msg = "Running `" + ' '.join(args) + "`"
         if kwargs["strip"] and not kwargs["log"] is True or kwargs["log"] is False:
             logging.debug(msg)
@@ -110,23 +118,34 @@ class Shell(object):
             stdin=sys.stdin
             stderr=sys.stderr
         else:
-            stdout=subprocess.PIPE
-            stdin=subprocess.PIPE
-            stderr=subprocess.PIPE
+            stdout=kwargs["stdout"]
+            stdin=kwargs["stdin"]
+            stderr=kwargs["stderr"]
+        env = None
+        if kwargs["addenv"]:
+            env = dict(os.environ.items() + kwargs["addenv"].items())
         # XXX: There is a possible problem here where we can fill up
         # the kernel buffer if we have 64KB of data.  This shouldn't
-        # be a problem, and the fix for such case would be to write to
+        # normally be a problem, and the fix for such case would be to write to
         # temporary files instead of a pipe.
+        #
+        # However, it *is* a problem when you do something silly, like
+        # pass --debug to mass-upgrade.
+        #
         # Another possible way of fixing this is converting from a
         # waitpid() pump to a select() pump, creating a pipe to
-        # ourself, and then setting up a
-        # SIGCHILD handler to write a single byte to the pipe to get
-        # us out of select() when a subprocess exits.
-        proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin)
-        if hasattr(self, "_async"):
-            self._async(proc, args, **kwargs)
+        # ourself, and then setting up a SIGCHILD handler to write a single
+        # byte to the pipe to get us out of select() when a subprocess exits.
+        proc = subprocess.Popen(args, stdout=stdout, stderr=stderr, stdin=stdin, cwd=self.cwd, env=env)
+        if self._async(proc, args, **kwargs):
             return proc
         stdout, stderr = proc.communicate(kwargs["input"])
+        # can occur if we were doing interactive communication; i.e.
+        # we didn't pass in PIPE.
+        if stdout is None:
+            stdout = ""
+        if stderr is None:
+            stderr = ""
         if not kwargs["interactive"]:
             if kwargs["strip"]:
                 self._log(None, stderr)
@@ -137,7 +156,7 @@ class Shell(object):
             else: eclass = CallError
             raise eclass(proc.returncode, args, stdout, stderr)
         if kwargs["strip"]:
-            return stdout.rstrip("\n")
+            return str(stdout).rstrip("\n")
         return (stdout, stderr)
     def _log(self, stdout, stderr):
         """Logs the standard output and standard input from a command."""
@@ -145,6 +164,10 @@ class Shell(object):
             logging.debug("STDOUT:\n" + stdout)
         if stderr:
             logging.debug("STDERR:\n" + stderr)
+    def _wait(self):
+        pass
+    def _async(self, *args, **kwargs):
+        return False
     def callAsUser(self, *args, **kwargs):
         """
         Performs a system call as a different user.  This is only possible
@@ -180,6 +203,8 @@ class Shell(object):
         on working directory context.  Keyword arguments are the
         same as :meth:`call`.
         """
+        if os.getuid():
+            return self.call(*args, **kwargs)
         uid = os.stat(os.getcwd()).st_uid
         # consider also checking ruid?
         if uid != os.geteuid():
@@ -199,6 +224,25 @@ class Shell(object):
         """
         kwargs["strip"] = True
         return self.call(*args, **kwargs)
+    def setcwd(self, cwd):
+        """
+        Sets the directory processes are executed in. This sets a value
+        to be passed as the ``cwd`` argument to ``subprocess.Popen``.
+        """
+        self.cwd = cwd
+    def interactive():
+        user_shell = os.getenv("SHELL")
+        if not user_shell: user_shell = "/bin/bash"
+        # XXX: scripts specific hack, since mbash doesn't respect the current working directory
+        # When the revolution comes (i.e. $ATHENA_HOMEDIR/Scripts is your Scripts home
+        # directory) this isn't strictly necessary, but we'll probably need to support
+        # web_scripts directories ad infinitum.
+        if user_shell == "/usr/local/bin/mbash": user_shell = "/bin/bash"
+
+        try:
+            self.call(user_shell, "-i", interactive=True)
+        except shell.CallError as e:
+            logging.warning("Shell returned non-zero exit code %d" % e.code)
 
 class ParallelShell(Shell):
     """
@@ -259,12 +303,13 @@ class ParallelShell(Shell):
         execution.  See :meth:`Shell.call` source code for details.
         """
         self.running[proc.pid] = (proc, args, python, on_success, on_error)
+        return True # so that the parent function returns
     def _wait(self):
         """
         Blocking call that waits for an open subprocess slot.  This is
         automatically called by :meth:`Shell.call`.
         """
-        # XXX: This API sucks; the actuall call/callAsUser call should
+        # XXX: This API sucks; the actual call/callAsUser call should
         # probably block automatically (unless I have a good reason not to)
         # bail out immediately on initial ramp up
         if len(self.running) < self.max: return
@@ -297,7 +342,16 @@ class ParallelShell(Shell):
             on_error(eclass(proc.returncode, args, stdout, stderr))
             return
         on_success(stdout, stderr)
+    def interactive():
+        raise Error("Cannot use interactive() on parallel shell")
 
+# Setup a convenience global instance
+shell = Shell()
+call = shell.call
+callAsUser = shell.callAsUser
+safeCall = shell.safeCall
+eval = shell.eval
+interactive = shell.interactive
 
 class DummyParallelShell(ParallelShell):
     """Same API as :class:`ParallelShell`, but doesn't actually