Refactor interactive shell and add addenv support.

Signed-off-by: Edward Z. Yang <ezyang@cs.stanford.edu>

diff --git a/wizard/command/upgrade.py b/wizard/command/upgrade.py
index 0850994247275912d37487767d39553ec43a21c9..0001bc5b6c69707ee929b10f7603d6983893f97c 100644 (file)
--- a/wizard/command/upgrade.py
+++ b/wizard/command/upgrade.py
@@ -350,13 +350,6 @@ class Upgrade(object):
             print "%d %s" % (conflicts, self.temp_wc_dir)
             raise MergeFailed
         else:
-            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"
             while 1:
                 print
                 print "ERROR: The merge failed with %d conflicts in these files:" % conflicts
@@ -371,10 +364,7 @@ class Upgrade(object):
                 print "NOTE: If you resolve these conflicts, and then the upgrade fails for"
                 print "an unrelated reason, you can run 'wizard upgrade --continue' from this"
                 print "directory to try again."
-                try:
-                    shell.call(user_shell, "-i", interactive=True)
-                except shell.CallError as e:
-                    logging.warning("Shell returned non-zero exit code %d" % e.code)
+                shell.interactive()
                 if shell.eval("git", "ls-files", "--unmerged").strip():
                     print
                     print "WARNING: There are still unmerged files."

diff --git a/wizard/shell.py b/wizard/shell.py
index 2197cdecc3cf219beebe5d5afe10744917682869..523dea23af6e1e7d1887b2777aceae6e8c1ca83c 100644 (file)
--- a/wizard/shell.py
+++ b/wizard/shell.py
@@ -76,6 +76,7 @@ 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.
@@ -96,6 +97,7 @@ class Shell(object):
         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)
@@ -119,6 +121,9 @@ class Shell(object):
             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
         # normally be a problem, and the fix for such case would be to write to
@@ -131,7 +136,7 @@ class Shell(object):
         # 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, cwd=self.cwd, )
+        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"])
@@ -225,6 +230,19 @@ class Shell(object):
         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):
     """
@@ -324,6 +342,8 @@ 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()
@@ -331,6 +351,7 @@ 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