[recipes] Add env recipe module

Simplifies the process of nesting environments. Usage:

with api.env({'myvar': 'myval'}):
  # do stuff

Same as api.step.context but specialized for just environment and merges
PATH variable intelligently.

Bug: skia:6473
Change-Id: I5769c69cbbbcdab0c6298cee6c5e1fe9caf89c78
Reviewed-on: https://skia-review.googlesource.com/14189
Reviewed-by: Kevin Lubick <kjlubick@google.com>
Commit-Queue: Eric Boren <borenet@google.com>

8 files changed, 113 insertions, 23 deletions

diff --git a/infra/bots/recipe_modules/env/__init__.py b/infra/bots/recipe_modules/env/__init__.py
new file mode 100644
index 0000000000..67d72f4fb8
--- /dev/null
+++ b/infra/bots/recipe_modules/env/__init__.py
@@ -0,0 +1,9 @@
+# Copyright 2017 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+DEPS = [
+  'recipe_engine/step',
+]
+
diff --git a/infra/bots/recipe_modules/env/api.py b/infra/bots/recipe_modules/env/api.py
new file mode 100644
index 0000000000..8060a0f946
--- /dev/null
+++ b/infra/bots/recipe_modules/env/api.py
@@ -0,0 +1,20 @@
+# Copyright 2017 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+from recipe_engine import recipe_api
+
+
+class EnvApi(recipe_api.RecipeApi):
+  def __call__(self, env_dict):
+    env = self.m.step.get_from_context('env', {})
+    # If PATH is defined in both, merge them together, merging default_env into
+    # path by replacing %(PATH)s
+    upstream_path = env.get('PATH', '')
+    env.update(env_dict)
+    my_path = env_dict.get('PATH', '')
+    if upstream_path and my_path and upstream_path != my_path:
+      env['PATH'] = upstream_path.replace(r'%(PATH)s', my_path)
+
+    return self.m.step.context({'env': env})
diff --git a/infra/bots/recipe_modules/env/example.expected/test.json b/infra/bots/recipe_modules/env/example.expected/test.json
new file mode 100644
index 0000000000..be38b1542c
--- /dev/null
+++ b/infra/bots/recipe_modules/env/example.expected/test.json
@@ -0,0 +1,44 @@
+[
+  {
+    "cmd": [
+      "echo",
+      "hi"
+    ],
+    "name": "1"
+  },
+  {
+    "cmd": [
+      "echo",
+      "hi"
+    ],
+    "env": {
+      "MYVAR": "myval"
+    },
+    "name": "2"
+  },
+  {
+    "cmd": [
+      "echo",
+      "hi"
+    ],
+    "env": {
+      "PATH": "mypath:%(PATH)s"
+    },
+    "name": "3"
+  },
+  {
+    "cmd": [
+      "echo",
+      "hi"
+    ],
+    "env": {
+      "PATH": "mypath:%(PATH)s:otherpath"
+    },
+    "name": "4"
+  },
+  {
+    "name": "$result",
+    "recipe_result": null,
+    "status_code": 0
+  }
+]
\ No newline at end of file
diff --git a/infra/bots/recipe_modules/env/example.py b/infra/bots/recipe_modules/env/example.py
new file mode 100644
index 0000000000..a22899b576
--- /dev/null
+++ b/infra/bots/recipe_modules/env/example.py
@@ -0,0 +1,25 @@
+# Copyright 2017 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+DEPS = [
+  'env',
+  'recipe_engine/step',
+]
+
+
+def RunSteps(api):
+  api.step('1', cmd=['echo', 'hi'])
+  with api.env({'MYVAR': 'myval'}):
+    api.step('2', cmd=['echo', 'hi'])
+
+  path = 'mypath:%(PATH)s'
+  with api.step.context({'env': {'PATH': path}}):
+    api.step('3', cmd=['echo', 'hi'])
+    with api.env({'PATH': '%(PATH)s:otherpath'}):
+      api.step('4', cmd=['echo', 'hi'])
+
+
+def GenTests(api):
+  yield api.test('test')
diff --git a/infra/bots/recipe_modules/run/__init__.py b/infra/bots/recipe_modules/run/__init__.py
index cb4c1ae0ee..25e4c0dad5 100644
--- a/infra/bots/recipe_modules/run/__init__.py
+++ b/infra/bots/recipe_modules/run/__init__.py
@@ -4,6 +4,7 @@
 
 DEPS = [
   'build/file',
+  'env',
   'recipe_engine/json',
   'recipe_engine/path',
   'recipe_engine/platform',
diff --git a/infra/bots/recipe_modules/run/api.py b/infra/bots/recipe_modules/run/api.py
index 87700d1f54..b487a75d9c 100644
--- a/infra/bots/recipe_modules/run/api.py
+++ b/infra/bots/recipe_modules/run/api.py
@@ -70,10 +70,9 @@ class SkiaStepApi(recipe_api.RecipeApi):
 
   def rmtree(self, path):
     """Wrapper around api.file.rmtree with environment fix."""
-    env = self.m.step.get_from_context('env', {})
-    env['PYTHONPATH'] = str(self.m.path['start_dir'].join(
-        'skia', 'infra', 'bots', '.recipe_deps', 'build', 'scripts'))
-    with self.m.step.context({'env': env}):
+    env = {'PYTHONPATH': str(self.m.path['start_dir'].join(
+        'skia', 'infra', 'bots', '.recipe_deps', 'build', 'scripts'))}
+    with self.m.env(env):
       self.m.file.rmtree(self.m.path.basename(path),
                          path,
                          infra_step=True)
@@ -81,18 +80,8 @@ class SkiaStepApi(recipe_api.RecipeApi):
   def __call__(self, steptype, name, abort_on_failure=True,
                fail_build_on_failure=True, **kwargs):
     """Run a step. If it fails, keep going but mark the build status failed."""
-    env = self.m.step.get_from_context('env', {})
-    # If PATH is defined in both, merge them together, merging default_env into
-    # path by replacing %(PATH)s
-    path = env.get('PATH', '')
-    env.update(self.m.vars.default_env)
-    default_path = self.m.vars.default_env.get('PATH', '')
-    if path and default_path and path != default_path:
-      path = path.replace(r'%(PATH)s', default_path)
-      env['PATH'] = path
-
     try:
-      with self.m.step.context({'env': env}):
+      with self.m.env(self.m.vars.default_env):
         return steptype(name=name, **kwargs)
     except self.m.step.StepFailure as e:
       if abort_on_failure or fail_build_on_failure:
diff --git a/infra/bots/recipes/perf.py b/infra/bots/recipes/perf.py
index 33bfba6d83..57ce23260b 100644
--- a/infra/bots/recipes/perf.py
+++ b/infra/bots/recipes/perf.py
@@ -12,6 +12,8 @@ import calendar
 DEPS = [
   'build/file',
   'core',
+  'env',
+  'flavor',
   'recipe_engine/json',
   'recipe_engine/path',
   'recipe_engine/platform',
@@ -20,7 +22,6 @@ DEPS = [
   'recipe_engine/step',
   'recipe_engine/time',
   'run',
-  'flavor',
   'vars',
 ]
 
@@ -254,7 +255,7 @@ def perf_steps(api):
       if not k in keys_blacklist:
         args.extend([k, api.vars.builder_cfg[k]])
 
-  env = api.step.get_from_context('env', {})
+  env = {}
   if 'Ubuntu16' in api.vars.builder_name:
     # The vulkan in this asset name simply means that the graphics driver
     # supports Vulkan. It is also the driver used for GL code.
@@ -282,7 +283,7 @@ def perf_steps(api):
   if '_AbandonGpuContext' in api.vars.builder_cfg.get('extra_config', ''):
     args.extend(['--abandonGpuContext', '--nocpu'])
 
-  with api.step.context({'env': env}):
+  with api.env(env):
     api.run(api.flavor.step, target, cmd=args,
             abort_on_failure=False)
 
@@ -296,10 +297,10 @@ def perf_steps(api):
 
 def RunSteps(api):
   api.core.setup()
-  env = api.step.get_from_context('env', {})
+  env = {}
   if 'iOS' in api.vars.builder_name:
     env['IOS_BUNDLE_ID'] = 'com.google.nanobench'
-  with api.step.context({'env': env}):
+  with api.env(env):
     try:
       if 'Chromecast' in api.vars.builder_name:
         api.flavor.install(resources=True, skps=True)
diff --git a/infra/bots/recipes/test.py b/infra/bots/recipes/test.py
index 5a21bd6eb1..8625b11bcb 100644
--- a/infra/bots/recipes/test.py
+++ b/infra/bots/recipes/test.py
@@ -9,6 +9,8 @@
 DEPS = [
   'build/file',
   'core',
+  'env',
+  'flavor',
   'recipe_engine/json',
   'recipe_engine/path',
   'recipe_engine/platform',
@@ -16,7 +18,6 @@ DEPS = [
   'recipe_engine/python',
   'recipe_engine/raw_io',
   'recipe_engine/step',
-  'flavor',
   'run',
   'vars',
 ]
@@ -647,7 +648,7 @@ def test_steps(api):
     args.append(skip_flag)
   args.extend(dm_flags(api.vars.builder_name))
 
-  env = api.step.get_from_context('env', {})
+  env = {}
   if 'Ubuntu16' in api.vars.builder_name:
     # The vulkan in this asset name simply means that the graphics driver
     # supports Vulkan. It is also the driver used for GL code.
@@ -677,7 +678,7 @@ def test_steps(api):
   if '_PreAbandonGpuContext' in api.vars.builder_cfg.get('extra_config', ''):
     args.append('--preAbandonGpuContext')
 
-  with api.step.context({'env': env}):
+  with api.env(env):
     api.run(api.flavor.step, 'dm', cmd=args, abort_on_failure=False)
 
   if api.vars.upload_dm_results: