src/tools/swarming_client/tests/test_utils.py

   1 # Copyright 2014 The Swarming Authors. All rights reserved.
   2 # Use of this source code is governed under the Apache License, Version 2.0 that
   3 # can be found in the LICENSE file.
   4
   5 import os
   6
   7 _UMASK = None
   8
   9 class EnvVars(object):
  10   """Context manager for environment variables.
  11
  12   Passed a dict to the constructor it sets variables named with the key to the
  13   value.  Exiting the context causes all the variables named with the key to be
  14   restored to their value before entering the context.
  15   """
  16   def __init__(self, var_map):
  17     self.var_map = var_map
  18     self._backup = None
  19
  20   def __enter__(self):
  21     self._backup = os.environ
  22     os.environ = os.environ.copy()
  23     os.environ.update(self.var_map)
  24
  25   def __exit__(self, exc_type, exc_value, traceback):
  26     os.environ = self._backup
  27
  28
  29 def umask():
  30   """Returns current process umask without modifying it."""
  31   global _UMASK
  32   if _UMASK is None:
  33     _UMASK = os.umask(0777)
  34     os.umask(_UMASK)
  35   return _UMASK