[lit] implement 'env' as shell builtin

[• Tzafrir]

Summary

Add implementation of env command as an internal shell builtin.

[rnk]

This has a few bugs. 'env' should only affect a single command, and it should only affect one command in the pipeline. The env should be reset for the next command in the pipeline. I'm drafting a patch to handle these edge cases, and I'll let you know when it lands.

llvm/trunk/utils/lit/lit/TestRunner.py

class ShellEnvironment(object):
 
     """Mutable shell environment containing things like CWD and env vars.
 
     Environment variables are not implemented, but cwd tracking is.
     """
 
     def __init__(self, cwd, env):
         self.cwd = cwd
-        self.env = env
+        self.env = dict(env)
 
 def executeShCmd(cmd, shenv, results):
     if isinstance(cmd, ShUtil.Seq):
         if cmd.op == ';':
             res = executeShCmd(cmd.lhs, shenv, results)
             return executeShCmd(cmd.rhs, shenv, results)
 
         if cmd.op == '&':
if isinstance(cmd, ShUtil.Seq):
                 res = executeShCmd(cmd.rhs, shenv, results)
             return res
 
         raise ValueError('Unknown shell command: %r' % cmd.op)
     assert isinstance(cmd, ShUtil.Pipeline)
 
     # Handle shell builtins first.
     if cmd.commands[0].args[0] == 'cd':
-        # Update the cwd in the environment.
+        if len(cmd.commands) != 1:
+            raise ValueError("'cd' cannot be part of a pipeline")
         if len(cmd.commands[0].args) != 2:
-            raise ValueError('cd supports only one argument')
+            raise ValueError("'cd' supports only one argument")
         newdir = cmd.commands[0].args[1]
+        # Update the cwd in the parent environment.
         if os.path.isabs(newdir):
             shenv.cwd = newdir
         else:
             shenv.cwd = os.path.join(shenv.cwd, newdir)
+        # The cd builtin always succeeds. If the directory does not exist, the
+        # following Popen calls will fail instead.
         return 0
 
     procs = []
     input = subprocess.PIPE
     stderrTempFiles = []
     opened_files = []
     named_temp_files = []
     # To avoid deadlock, we use a single stderr stream for piped
     # output. This is null until we have seen some output using
     # stderr.
     for i,j in enumerate(cmd.commands):
+        # Reference the global environment by default.
+        cmd_shenv = shenv
+        if j.args[0] == 'env':
+            # Create a copy of the global environment and modify it for this one
+            # command. There might be multiple envs in a pipeline:
+            #   env FOO=1 llc < %s | env BAR=2 llvm-mc | FileCheck %s
+            cmd_shenv = ShellEnvironment(shenv.cwd, shenv.env)
+            arg_idx = 1
+            for arg_idx, arg in enumerate(j.args[1:]):
+                # Partition the string into KEY=VALUE.
+                key, eq, val = arg.partition('=')
+                # Stop if there was no equals.
+                if eq == '':
+                    break
+                cmd_shenv.env[key] = val
+            j.args = j.args[arg_idx+1:]
+
         # Apply the redirections, we use (N,) as a sentinel to indicate stdin,
         # stdout, stderr for N equal to 0, 1, or 2 respectively. Redirects to or
         # from a file are represented with a list [file, mode, file-object]
         # where file-object is initially None.
         redirects = [(0,), (1,), (2,)]
         for r in j.redirects:
             if r[0] == ('>',2):
                 redirects[2] = [r[1], 'w', None]
for i,j in enumerate(cmd.commands):
                     raise InternalShellError(j,"Unsupported redirect on stdout")
                 result = subprocess.PIPE
             else:
                 if r[2] is None:
                     if kAvoidDevNull and r[0] == '/dev/null':
                         r[2] = tempfile.TemporaryFile(mode=r[1])
                     else:
                         # Make sure relative paths are relative to the cwd.
-                        redir_filename = os.path.join(shenv.cwd, r[0])
+                        redir_filename = os.path.join(cmd_shenv.cwd, r[0])
                         r[2] = open(redir_filename, r[1])
                     # Workaround a Win32 and/or subprocess bug when appending.
                     #
                     # FIXME: Actually, this is probably an instance of PR6753.
                     if r[1] == 'a':
                         r[2].seek(0, 2)
                     opened_files.append(r[2])
                 result = r[2]
for i,j in enumerate(cmd.commands):
             #
             # FIXME: This is slow, but so is deadlock.
             if stderr == subprocess.PIPE and j != cmd.commands[-1]:
                 stderr = tempfile.TemporaryFile(mode='w+b')
                 stderrTempFiles.append((i, stderr))
 
         # Resolve the executable path ourselves.
         args = list(j.args)
-        executable = lit.util.which(args[0], shenv.env['PATH'])
+        executable = lit.util.which(args[0], cmd_shenv.env['PATH'])
         if not executable:
             raise InternalShellError(j, '%r: command not found' % j.args[0])
 
         # Replace uses of /dev/null with temporary files.
         if kAvoidDevNull:
             for i,arg in enumerate(args):
                 if arg == "/dev/null":
                     f = tempfile.NamedTemporaryFile(delete=False)
                     f.close()
                     named_temp_files.append(f.name)
                     args[i] = f.name
 
         try:
-            procs.append(subprocess.Popen(args, cwd=shenv.cwd,
+            procs.append(subprocess.Popen(args, cwd=cmd_shenv.cwd,
                                           executable = executable,
                                           stdin = stdin,
                                           stdout = stdout,
                                           stderr = stderr,
-                                          env = shenv.env,
+                                          env = cmd_shenv.env,
                                           close_fds = kUseCloseFDs))
         except OSError as e:
             raise InternalShellError(j, 'Could not create process due to {}'.format(e))
 
         # Immediately close stdin for any process taking stdin from us.
         if stdin == subprocess.PIPE:
             procs[-1].stdin.close()
             procs[-1].stdin = None