[lit] add support for unsetting environment variable

[jlpeyton]

Summary

This patch adds support for the -u flag for the env command. The -u flag allows users to explicitly unset an environment variable for a particular command. This is useful for OpenMP which heavily relies on environment variables to determine the library's behavior in certain situations. So being able to unset a potentially polluted environment is important.

Example usage:
env -u FOO -u BAR GOO=16 CAR=5
This removes both FOO and BAR from the environment while setting GOO=16 and HAR=5.

[jlpeyton]

ping

[• ddunbar]

This looks reasonable to me, can we get a test for the behavior though?

[jlpeyton]

Added testing for env and env -u feature of ShTest format

[jlpeyton]

ping

[• ddunbar]

LGTM

[jlpeyton]

D36403 committed this functionality (a74a4df4d5 or svn r311180) . This revision is not needed.

lit/TestRunner.py

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
+            unset_next_env_var = False
             for arg_idx, arg in enumerate(j.args[1:]):
+                # Support for the -u flag (unsetting) for env command
+                # e.g., env -u FOO -u BAR will remove both FOO and BAR
+                # from the environment.
+                if arg == '-u':
+                    unset_next_env_var = True
+                    continue
+                if unset_next_env_var:
+                    unset_next_env_var = False
+                    if arg in cmd_shenv.env:
+                        del cmd_shenv.env[arg]
+                    continue
+
                 # 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:]
 

tests/Inputs/shtest-env/env-u.txt

+# Check and make sure preset environment variable were set in lit.cfg
+#
+# RUN: %{python} print_environment \
+# RUN:   | FileCheck --check-prefix=CHECK-ENV-PRESET %s
+#
+# Check single unset of environment variable
+#
+# RUN: env -u FOO %{python} print_environment \
+# RUN:  | FileCheck --check-prefix=CHECK-ENV-UNSET-1 %s
+#
+# Check multiple unsets of environment variables
+#
+# RUN: env -u FOO -u BAR %{python} print_environment \
+# RUN:  | FileCheck --check-prefix=CHECK-ENV-UNSET-MULTIPLE %s
+
+# CHECK-ENV-PRESET: BAR = 2
+# CHECK-ENV-PRESET: FOO = 1
+
+# CHECK-ENV-UNSET-1: BAR = 2
+# CHECK-ENV-UNSET-1-NOT: FOO
+
+# CHECK-ENV-UNSET-MULTIPLE-NOT: BAR
+# CHECK-ENV-UNSET-MULTIPLE-NOT: FOO

tests/Inputs/shtest-env/env.txt

+# Check for simple one environment variable setting
+#
+# RUN: env A_FOO=999 %{python} print_environment \
+# RUN:   | FileCheck --check-prefix=CHECK-ENV-1 %s
+#
+# Check for multiple environment variable settings
+#
+# RUN: env A_FOO=1 B_BAR=2 C_OOF=3 %{python} print_environment \
+# RUN:   | FileCheck --check-prefix=CHECK-ENV-MULTIPLE %s
+
+# CHECK-ENV-1: A_FOO = 999
+
+# CHECK-ENV-MULTIPLE: A_FOO = 1
+# CHECK-ENV-MULTIPLE: B_BAR = 2
+# CHECK-ENV-MULTIPLE: C_OOF = 3

tests/Inputs/shtest-env/lit.cfg

+import lit.formats
+config.name = 'shtest-env'
+config.suffixes = ['.txt']
+config.test_format = lit.formats.ShTest()
+config.test_source_root = None
+config.test_exec_root = None
+config.environment['FOO'] = '1'
+config.environment['BAR'] = '2'
+config.substitutions.append(('%{python}', sys.executable))

tests/Inputs/shtest-env/mixed.txt

+# Check for setting and removing one environment variable
+#
+# RUN: env A_FOO=999 -u FOO %{python} print_environment \
+# RUN:   | FileCheck --check-prefix=CHECK-ENV-1 %s
+#
+# Check for setting/unsetting multiple environment variables
+#
+# RUN: env A_FOO=1 -u FOO B_BAR=2 -u BAR C_OOF=3 %{python} print_environment \
+# RUN:   | FileCheck --check-prefix=CHECK-ENV-MULTIPLE %s
+
+# CHECK-ENV-1: A_FOO = 999
+# CHECK-ENV-1-NOT: FOO
+
+# CHECK-ENV-MULTIPLE: A_FOO = 1
+# CHECK-ENV-MULTIPLE-NOT: BAR
+# CHECK-ENV-MULTIPLE: B_BAR = 2
+# CHECK-ENV-MULTIPLE: C_OOF = 3
+# CHECK-ENV-MULTIPLE-NOT: FOO

tests/Inputs/shtest-env/print_environment

+#!/usr/bin/env python
+
+import os
+
+sorted_environment = sorted(os.environ.items())
+
+for name,value in sorted_environment:
+    print name,'=',value

tests/shtest-env.py

+# Check the env command
+#
+# RUN: %{lit} -a -v %{inputs}/shtest-env