[context] Add path prefix/suffix manipulation.

recipe_modules/context/api.py

 # Copyright 2017 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 """The context module provides APIs for manipulating a few pieces of 'ambient'
 data that affect how steps are run:
   cwd - The current working directory.
   env - The environment variables.
   infra_step - Whether or not failures should be treated as infrastructure
     failures vs. normal failures.
   name_prefix - A prefix for all step names.
   nest_level - An indicator for the UI of how deeply to nest steps.
 
 The values here are all scoped using Python's `with` statement; there's no
 mechanism to make an open-ended adjustment to these values (i.e. there's no way
 to change the cwd permanently for a recipe, except by surrounding the entire
 recipe with a with statement). This is done to avoid the surprises that
 typically arise with things like os.environ or os.chdir in a normal python
 program.
 
 Example:
   with api.context(cwd=api.path['start_dir'].join('subdir')):
     # this step is run inside of the subdir directory.
     api.step("cat subdir/foo", ['cat', './foo'])
 """
 
 
 import collections
+import os
 
 from contextlib import contextmanager
 
 from recipe_engine import recipe_api
 from recipe_engine.config_types import Path
 from recipe_engine.recipe_api import RecipeApi
 
 
 def check_type(name, var, expect):
   if not isinstance(var, expect):  # pragma: no cover
     raise TypeError('%s is not %s: %r (%s)' % (
       name, expect.__name__, var, type(var).__name__))
 
 
 class ContextApi(RecipeApi):
+
   # TODO(iannucci): move implementation of these data directly into this class.
   def __init__(self, **kwargs):
     super(RecipeApi, self).__init__(**kwargs)
 
     self._cwd = [None]
     self._env = [{}]
     self._infra_step = [False]
     self._name_prefix = ['']
     # this could be a number, but it makes the logic easier to use a stack.
     self._nest_level = [0]
+    self._prefix_paths = [{}]
+    self._suffix_paths = [{}]
 
   @contextmanager
   def __call__(self, cwd=None, env=None, increment_nest_level=None,
-               infra_steps=None, name_prefix=None):
+               infra_steps=None, name_prefix=None, path_prefix=None,
+               path_suffix=None):
     """Allows adjustment of multiple context values in a single call.
 
     Contextual data:
       * cwd (Path) - the current working directory to use for all steps.
         To 'reset' to the original cwd at the time recipes started, pass
         `api.path['start_dir']`.
       * infra_steps (bool) - if steps in this context should be considered
         infrastructure steps. On failure, these will raise InfraFailure
         exceptions instead of StepFailure exceptions.
       * increment_nest_level (True) - increment the nest level by 1 in this
         context. Typically you won't directly interact with this, but should
         use api.step.nest instead.
       * name_prefix (str) - A string to prepend to the names of all steps in
         this context. These compose with '.' characters if multiple name prefix
         contexts occur. See below for more info.
       * env (dict) - Environmental variable overrides. See below for more info.

[iannucci, 2017/06/05 19:23:04]

wdyt about something like:

  with api.context(env={"PATH": api.context.Prefix(a, b, c)}):

Then we don't need to additional kwargs, just Prefix and Suffix. I would also
give Prefix and Suffix an __add__ so you could do `Prefix(...)+Suffix(...)`. I
would strongly consider banning mixed use (i.e. an envvar is either
Prefix/Suffix, or it's a string, not both).

[dnj, 2017/06/07 02:59:29]

That seems fine. For now I think we need mixed-use to transition to the new
notation, but I'll add a TODO to forbid mixed-use once that's done.

+      * path_prefix (dict) - Environmental variable prefix path overrides.
+        Consider using the "path_prefix" and "path_remove" helper functions
+        instead of directly manipulating this value.
+      * path_suffix (dict) - Environmental variable suffix path overrides.
+        Consider using the "path_suffix" and "path_remove" helper functions
+        instead of directly manipulating this value.

[iannucci, 2017/06/05 19:23:04]

I don't like the helper function idea; we should keep __call__ as the only way
to set stuff in the context module. I have a plan to clean up the api.step.nest
thing later too, in case that's what you're modeling this after.

[dnj, 2017/06/07 02:59:29]

SGTM

 
     Name prefixes:
 
     Multiple invocations concatenate values with '.'.
 
     Example:
       with api.context(name_prefix='hello'):
         # has name 'hello.something'
         api.step('something', ['echo', 'something'])
 
         with api.context(name_prefix='world'):
           # has name 'hello.world.other'
           api.step('other', ['echo', 'other'])
 
     Environmental Variable Overrides:
 
     Env is a mapping of environment variable name to the value you want that
     environment variable to have. The value is a string, with a couple
     exceptions:
       * If value is None, this environment variable will be removed from the
         environment when the step runs.
       * String values will be %-formatted with the current value of the
         environment at the time the step runs. This means that you can have
         a value like:
           "/path/to/my/stuff:%(PATH)s"
         Which, at the time the step executes, will inject the current value of
         $PATH.
 
-    TODO(iannucci): implement env_paths which allows for easier manipulation of
-    `pathsep` environment variables like $PATH, $PYTHONPATH, etc.
+    path_prefix and path_suffix are mappings of environment variables to lists
+    of Path elements, meant to be applied to "pathsep"-separated environment
+    variables such as PATH and PYTHONPATH. These values are prefixed and/or
+    appended to their environment variable key, separated by the OS-specific
+    path separator.
+      * If key begins with a "!" (e.g., "!PATH"), the named element will be
+        removed from its environment instead of added.
+      * Consider using helper functions instead of directly manipulating this
+        value.
 
     TODO(iannucci): combine nest_level and name_prefix
 
     Example:
       # suppose the OS's envar $OTHER is set to "yes"
       with api.context(env={'ENV_VAR': 'something:%(OTHER)s'}):
         # environment updates are additive.
         with api.context(env={'OTHER': 'cool:%(OTHER)s'}):
           # echos 'something:yes'
           # Note that the substitution always happens with the system
@@ skipping 53 matching lines @@
             # will raise an exception. If not, then this is a pluasible format
             # string.
             ('%(foo)s'+v) % collections.defaultdict(str)
           except Exception:
             raise ValueError(('Invalid %%-formatting parameter in envvar, '
                               'only %%(ENVVAR)s allowed: %r') % (v,))
         new[k] = v
       self._env.append(new)
       to_pop.append(self._env)
 
+    for var, collection, name, is_prefix in (
+        (path_prefix, self._prefix_paths, 'path_prefix', True),
+        (path_suffix, self._suffix_paths, 'path_suffix', False)):
+      if var is None or var == {}:
+        continue
+      check_type(name, var, dict)
+      new = dict(collection[-1])
+      for k, paths in var.iteritems():
+        check_type(k, paths, collections.Iterable)

[iannucci, 2017/06/05 19:23:04]

check_type("KEY %s[%r]" % (name, k), ...)

so then the error looks like:

  TypeError('KEY path_prefix["foobar"] is not ....')

[dnj, 2017/06/07 02:59:28]

No longer needed.

+        remove = k.startswith('!')
+        if remove:
+          k = k[1:]

[iannucci, 2017/06/05 19:23:04]

let's remove the removal functionality for now

[dnj, 2017/06/07 02:59:29]

Done.

+
+        # Remove any occurrences of the current value and any duplicates from
+        # "paths". This will keep the prefix/suffix lists clean.
+        nv = new.get(k, ())
+        seen = set()
+        unique_paths = []
+        for i, path in enumerate(paths):
+          check_type('%s element #%d' % (k, i), path, Path)

[iannucci, 2017/06/05 19:23:04]

'%s[%r]' % (name, k)

[dnj, 2017/06/07 02:59:28]

Also no longer needed.

+          nv = tuple(x for x in nv if x != path)
+          if path not in seen:
+            unique_paths.append(path)
+            seen.add(path)
+        if not remove:
+          if is_prefix:
+            nv = tuple(unique_paths) + nv
+          else:
+            nv += tuple(unique_paths)
+        if nv:
+          new[k] = nv
+        else:
+          new.pop(k, None)
+      collection.append(new)
+      to_pop.append(collection)
+
     try:
       yield
     finally:
       for p in to_pop:
         p.pop()
 
   @property
   def cwd(self):
     """Returns the current working directory that steps will run in.
 
     Returns (Path|None) - The current working directory. A value of None is
       equivalent to api.path['start_dir'], though only occurs if no cwd has been
       set (e.g. in the outermost context of RunSteps).
     """
     return self._cwd[-1]
 
   @property
   def env(self):
     """Returns modifications to the environment.
 
     By default this is empty; There's no facility to observe the program's
     startup environment. If you want to pass data to the recipe, it should be
     done with properties.
 
     Returns (dict) - The env-key -> value mapping of current environment
       modifications.
     """
     # TODO(iannucci): store env in an immutable way to avoid excessive copies.
     # TODO(iannucci): handle case-insensitive keys on windows
-    return dict(self._env[-1])
+    env = dict(self._env[-1])
+    for collection, is_prefix in (
+        (self._prefix_paths, True),
+        (self._suffix_paths, False)):
+      for key, paths in collection[-1].iteritems():
+        cur = env.get(key, '%%(%s)s' % (key,))
+        cur = cur.split(os.pathsep) if cur else []

[iannucci, 2017/06/05 19:24:04]

this will fail in production/simulation; use self.m.path.pathsep.

[dnj, 2017/06/07 02:59:29]

Done.

+        if is_prefix:
+          cur = list(str(p) for p in paths) + cur
+        else:
+          cur += list(str(p) for p in paths)
+        env[key] = os.pathsep.join(cur)
+    return env

[iannucci, 2017/06/05 19:23:04]

I'm not sure how I feel about this; I think that step.__call__ should be the
only thing assembling this, and that all other uses of env should just deal with
env. In fact, I would go so far as to make it an error to have the same envvar
key manipulated via env and also the prefix/suffix methods.

Then we should add a prefix_paths and suffix_paths @property to keep this
class's API uniform.

[dnj, 2017/06/07 02:59:29]

Done.

 
   @property
   def infra_step(self):
     """Returns the current value of the infra_step setting.
 
     Returns (bool) - True iff steps are currently considered infra steps.
     """
     return self._infra_step[-1]
 
   @property
   def name_prefix(self):
     """Gets the current step name prefix.
 
     Returns (str) - The string prefix that every step will have prepended to it.
     """
     return self._name_prefix[-1]
 
   @property
   def nest_level(self):
     """Returns the current 'nesting' level.
 
     Note: This api is low-level, and you should always prefer to use
     `api.step.nest`. This api is included for completeness and documentation
     purposes.
 
     Returns (int) - The current nesting level.
     """
     return self._nest_level[-1]
+
+  @contextmanager
+  def path_prefix(self, key, *paths):
+    """Adds specified paths to the beginning of an environment variable.
+
+    Each path in paths is added, in order, as a prefix to the environment
+    variable "key", delimited by the OS path separator. This can be used for
+    easy manipulation of path environment variables such as PATH and PYTHONPATH.
+
+    If "key" begins with an "!", the specified paths will be removed from the
+    prefix. The "path_remove" helper function may be a more elegant way to
+    remove a path.
+
+    Args:
+      key (str): The environment variable key.
+      paths (...Path): The list of paths to prefix.
+    """
+    check_type('key', key, str)
+    with self(path_prefix={key: paths}):
+      yield
+
+  @contextmanager
+  def path_suffix(self, key, *paths):
+    """Adds specified paths to the end of an environment variable.
+
+    Each path in paths is added, in order, as a suffix to the environment
+    variable "key", delimited by the OS path separator. This can be used for
+    easy manipulation of path environment variables such as PATH and PYTHONPATH.
+
+    If "key" begins with an "!", the specified paths will be removed from the
+    prefix. The "path_remove" helper function may be a more elegant way to
+    remove a path.
+
+    Args:
+      key (str): The environment variable key.
+      paths (...Path): The list of paths to suffix.
+    """
+    check_type('key', key, str)
+    with self(path_suffix={key: paths}):
+      yield
+
+  @contextmanager
+  def path_remove(self, key, *paths):

[iannucci, 2017/06/05 19:23:04]

let's remove this functionality for now. I can't think of any use that wouldn't
be a horrible encapsulation violation. It would be easy to add later if
something needed it. The additive-only API is easier to reason about.

[dnj, 2017/06/07 02:59:29]

Done.

+    """Removes specified paths from an environment variable's prefix/suffix.
+
+    Any path in "paths" will be removed from the environment's prefix and suffix
+    paths. If the path is not present, it will be ignored.
+
+    WARNING: this will NOT remove a path that was not, itself, added as a
+    prefix or suffix. For example, if a path is added to the environment as a
+    string, rather than through "path_prefix" or "path_suffix", it will not be
+    removed.
+
+    Args:
+      key (str): The environment variable key.
+      paths (...Path): The list of paths to suffix.
+    """
+    check_type('key', key, str)
+    if not key.startswith('!'):
+      key = '!'+key
+    with self(path_prefix={key: paths}, path_suffix={key: paths}):
+      yield