GR-16806: add GraalPythonVM env support

Author: cosminbasca

mx.graalpython/mx_graalpython_benchmark.py

@@ -33,6 +33,7 @@
 import mx_benchmark
 from mx_benchmark import StdOutRule, java_vm_registry, Vm, GuestVm, VmBenchmarkSuite, AveragingBenchmarkMixin
 from mx_graalpython_bench_param import HARNESS_PATH
+from contextlib import contextmanager
 
 # ----------------------------------------------------------------------------------------------------------------------
 #
@@ -75,6 +76,27 @@ def _check_vm_args(name, args):
                  "got {} instead".format(args))
 
 
+@contextmanager
+def environ(env):
+    def _handle_var((k, v)):
+        if v is None:
+            del os.environ[k]
+        else:
+            os.environ[k] = str(v)
+
+    if env:
+        prev_env = {v: os.getenv(v) for v in env}
+        map(_handle_var, env.items())
+    else:
+        prev_env = None
+
+    try:
+        yield
+    finally:
+        if prev_env:
+            map(_handle_var, prev_env.items())
+
+
 # ----------------------------------------------------------------------------------------------------------------------
 #
 # the vm definitions
@@ -194,14 +216,15 @@ def name(self):
 
 class GraalPythonVm(GuestVm):
     def __init__(self, config_name=CONFIGURATION_DEFAULT, distributions=None, cp_suffix=None, cp_prefix=None,
-                 host_vm=None, extra_vm_args=None, extra_polyglot_args=None):
+                 host_vm=None, extra_vm_args=None, extra_polyglot_args=None, env=None):
         super(GraalPythonVm, self).__init__(host_vm=host_vm)
         self._config_name = config_name
         self._distributions = distributions
         self._cp_suffix = cp_suffix
         self._cp_prefix = cp_prefix
         self._extra_vm_args = extra_vm_args
         self._extra_polyglot_args = extra_polyglot_args
+        self._env = env
 
     def hosting_registry(self):
         return java_vm_registry
@@ -240,10 +263,11 @@ def run(self, cwd, args):
         cmd = truffle_options + vm_args + extra_polyglot_args + args
 
         host_vm = self.host_vm()
-        if hasattr(host_vm, 'run_lang'):
-            return host_vm.run_lang('graalpython', extra_polyglot_args + args, cwd)
-        else:
-            return host_vm.run(cwd, cmd)
+        with environ(self._env):
+            if hasattr(host_vm, 'run_lang'):
+                return host_vm.run_lang('graalpython', extra_polyglot_args + args, cwd)
+            else:
+                return host_vm.run(cwd, cmd)
 
     def name(self):
         return VM_NAME_GRAALPYTHON
@@ -254,7 +278,8 @@ def config_name(self):
     def with_host_vm(self, host_vm):
         return self.__class__(config_name=self._config_name, distributions=self._distributions,
                               cp_suffix=self._cp_suffix, cp_prefix=self._cp_prefix, host_vm=host_vm,
-                              extra_vm_args=self._extra_vm_args, extra_polyglot_args=self._extra_polyglot_args)
+                              extra_vm_args=self._extra_vm_args, extra_polyglot_args=self._extra_polyglot_args,
+                              env=self._env)
 
 
 # ----------------------------------------------------------------------------------------------------------------------