Python contextutil.environment_as函数代码示例

def test_environment_negation():
  with temporary_file() as output:
    with environment_as(HORK = 'BORK'):
      with environment_as(HORK = None):
      # test that the variable is cleared
        subprocess.Popen([sys.executable, '-c', 'import os; print os.environ.has_key("HORK")'],
          stdout=output).wait()
        output.seek(0)
        assert output.read() == 'False\n'

  def setUpClass(cls):
    cls.origin = safe_mkdtemp()
    with pushd(cls.origin):
      subprocess.check_call(['git', 'init', '--bare'])

    cls.gitdir = safe_mkdtemp()
    cls.worktree = safe_mkdtemp()

    cls.readme_file = os.path.join(cls.worktree, 'README')

    with environment_as(GIT_DIR=cls.gitdir, GIT_WORK_TREE=cls.worktree):
      cls.init_repo('depot', cls.origin)

      touch(cls.readme_file)
      subprocess.check_call(['git', 'add', 'README'])
      subprocess.check_call(['git', 'commit', '-am', 'initial commit with decode -> \x81b'])
      subprocess.check_call(['git', 'tag', 'first'])
      subprocess.check_call(['git', 'push', '--tags', 'depot', 'master'])
      subprocess.check_call(['git', 'branch', '--set-upstream', 'master', 'depot/master'])

      with safe_open(cls.readme_file, 'w') as readme:
        readme.write('Hello World.')
      subprocess.check_call(['git', 'commit', '-am', 'Update README.'])

    cls.clone2 = safe_mkdtemp()
    with pushd(cls.clone2):
      cls.init_repo('origin', cls.origin)
      subprocess.check_call(['git', 'pull', '--tags', 'origin', 'master:master'])

      with safe_open(os.path.realpath('README'), 'a') as readme:
        readme.write('--')
      subprocess.check_call(['git', 'commit', '-am', 'Update README 2.'])
      subprocess.check_call(['git', 'push', '--tags', 'origin', 'master'])

    cls.git = Git(gitdir=cls.gitdir, worktree=cls.worktree)

 def build_egg(self, egg_root, target):
   """Build an egg containing the files at egg_root for the specified target.
   There must be an egg_root/setup.py file."""
   # TODO(Brian Wickman): Do a sanity check somewhere to ensure that
   # setuptools is on the path?
   args = [
     sys.executable,
     'setup.py', 'bdist_egg',
     '--dist-dir=dist',
     '--bdist-dir=build.%s' % target.name]
   with pushd(egg_root):
     print 'EggBuilder executing: %s' % ' '.join(args)
     with environment_as(PYTHONPATH = ':'.join(sys.path)):
       po = subprocess.Popen(args, stderr=subprocess.PIPE, stdout=subprocess.PIPE)
       rv = po.wait()
     eggs = os.path.abspath(os.path.join('dist', '*.egg'))
     eggs = glob.glob(eggs)
     if rv != 0 or len(eggs) != 1:
       comm = po.communicate()
       print >> sys.stderr, 'egg generation failed (return value=%d, num eggs=%d)' % (
         rv, len(eggs))
       print >> sys.stderr, 'STDOUT'
       print >> sys.stderr, comm[0]
       print >> sys.stderr, 'STDERR'
       print >> sys.stderr, comm[1]
       raise EggBuilder.EggBuildingException(
         'Generation of eggs failed for target = %s' % target)
     egg_path = eggs[0]
   return egg_path

 def _maybe_scrubbed_env(cls):
   for env_var in cls._SCRUBBED_ENV:
     value = os.getenv(env_var)
     if value:
       log.warn('Scrubbing {env_var}={value}'.format(env_var=env_var, value=value))
   with environment_as(**cls._SCRUBBED_ENV):
     yield

  def run(self):
    self._run_count += 1
    atexit.register(self.cleanup)

    if self.script_filename:
      os.unlink(self.script_filename)

    with temporary_file(cleanup=False) as fp:
      self.script_filename = fp.name
      fp.write(self.RUN_JOB_SCRIPT % {
        'filename': self.job_filename,
        'sandbox': self.sandbox,
        'root': self.tempdir,
        'task_id': self.task_id,
        'state_filename': self.state_filename,
        'success_rate': self.success_rate,
        'random_seed': self.random_seed + self._run_count,
        'extra_task_runner_args': self.extra_task_runner_args,
      })

    with environment_as(PYTHONPATH=os.pathsep.join(sys.path)):
      self.po = subprocess.Popen([sys.executable, self.script_filename],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      try:
        so, se = self.po.communicate()
      except OSError as e:
        if e.errno == errno.ECHILD:
          so = se = 'Killed'
        else:
          raise

    rc = self.po.returncode
    if rc != 0:
      if os.path.exists(self.job_filename):
        with open(self.job_filename) as fp:
          config = fp.read()
      else:
        config = 'Nonexistent!'
      if 'THERMOS_DEBUG' in os.environ:
        print("Runner failed!\n\n\nconfig:%s\n\n\nstdout:%s\n\n\nstderr:%s\n\n\n" % (
            config, so, se))

    try:
      with open(self.state_filename, 'r') as fp:
        self.state = thrift_deserialize(RunnerState(), fp.read())
    except Exception as e:
      if 'THERMOS_DEBUG' in os.environ:
        print('Failed to load Runner state: %s' % e, file=sys.stderr)
      self.state = RunnerState()

    try:
      self.reconstructed_state = CheckpointDispatcher.from_file(
          self.pathspec.getpath('runner_checkpoint'))
    except Exception as e:
      print('Failed to replay checkpoint: %s' % e, file=sys.stderr)
      self.reconstructed_state = None
    self.initialized = True
    return rc

 def _maybe_scrubbed_classpath(self):
   if self._scrub_classpath:
     classpath = os.getenv('CLASSPATH')
     if classpath:
       log.warn('Scrubbing CLASSPATH=%s' % classpath)
     with environment_as(CLASSPATH=None):
       yield
   else:
     yield

 def config(self, overrides=''):
   """Returns a config valid for the test build root."""
   if overrides:
     with temporary_file() as fp:
       fp.write(overrides)
       fp.close()
       with environment_as(PANTS_CONFIG_OVERRIDE=fp.name):
         return Config.load()
   else:
     return Config.load()

  def test(self):
    self.assertEqual(set(), self.git.changed_files())
    self.assertEqual(set(['README']), self.git.changed_files(from_commit='HEAD^'))

    tip_sha = self.git.commit_id
    self.assertTrue(tip_sha)

    self.assertTrue(tip_sha in self.git.changelog())

    self.assertTrue(self.git.tag_name.startswith('first-'))
    self.assertEqual('master', self.git.branch_name)

    def edit_readme():
      with open(self.readme_file, 'a') as readme:
        readme.write('More data.')

    edit_readme()
    with open(os.path.join(self.worktree, 'INSTALL'), 'w') as untracked:
      untracked.write('make install')
    self.assertEqual(set(['README']), self.git.changed_files())
    self.assertEqual(set(['README', 'INSTALL']), self.git.changed_files(include_untracked=True))

    try:
      # These changes should be rejected because our branch point from origin is 1 commit behind
      # the changes pushed there in clone 2.
      self.git.commit('API Changes.')
    except Scm.RemoteException:
      with environment_as(GIT_DIR=self.gitdir, GIT_WORK_TREE=self.worktree):
        subprocess.check_call(['git', 'reset', '--hard', 'depot/master'])
      self.git.refresh()
      edit_readme()

    self.git.commit('''API '"' " Changes.''')
    self.git.tag('second', message='''Tagged ' " Changes''')

    with temporary_dir() as clone:
      with pushd(clone):
        subprocess.check_call(['git', 'init'])
        subprocess.check_call(['git', 'remote', 'add', 'origin', self.origin])
        subprocess.check_call(['git', 'pull', '--tags', 'origin', 'master:master'])

        with open(os.path.realpath('README')) as readme:
          self.assertEqual('--More data.', readme.read())

        git = Git()

        # Check that we can pick up committed and uncommitted changes.
        with safe_open(os.path.realpath('CHANGES'), 'w') as changes:
          changes.write('none')
        subprocess.check_call(['git', 'add', 'CHANGES'])
        self.assertEqual(set(['README', 'CHANGES']), git.changed_files(from_commit='first'))

        self.assertEqual('master', git.branch_name)
        self.assertEqual('second', git.tag_name)

def safe_classpath(logger=None):
  """
    Yields to a block in an environment with no CLASSPATH.  This is useful to ensure hermetic java
    invocations.
  """
  classpath = os.getenv('CLASSPATH')
  if classpath:
    logger = logger or log.warn
    logger('Scrubbing CLASSPATH=%s' % classpath)
  with environment_as(CLASSPATH=None):
    yield

def test_pex_python_symlink():
  with temporary_dir() as td:
    with environment_as(HOME=td):
      symlink_path = os.path.join(td, 'python-symlink')
      os.symlink(sys.executable, symlink_path)
      pexrc_path = os.path.join(td, '.pexrc')
      with open(pexrc_path, 'w') as pexrc:
        pexrc.write("PEX_PYTHON=%s" % symlink_path)

      body = "print('Hello')"
      _, rc = run_simple_pex_test(body, coverage=True)
      assert rc == 0

 def do_test_jre_env_var(self, env_var, env_value, scrubbed=True):
   with self.jre(env_var=env_var) as jre:
     executor = SubprocessExecutor(Distribution(bin_path=jre))
     with environment_as(**{env_var: env_value}):
       self.assertEqual(env_value, os.getenv(env_var))
       process = executor.spawn(classpath=['dummy/classpath'],
                                main='dummy.main',
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
       stdout, _ = process.communicate()
       self.assertEqual(0, process.returncode)
       self.assertEqual('' if scrubbed else env_value, stdout.strip())

def test_override_single_variable():
  with temporary_file() as output:
    # test that the override takes place
    with environment_as(HORK = 'BORK'):
      subprocess.Popen([sys.executable, '-c', 'import os; print os.environ["HORK"]'],
        stdout=output).wait()
      output.seek(0)
      assert output.read() == 'BORK\n'

    # test that the variable is cleared
    with temporary_file() as new_output:
      subprocess.Popen([sys.executable, '-c', 'import os; print os.environ.has_key("HORK")'],
        stdout=new_output).wait()
      new_output.seek(0)
      assert new_output.read() == 'False\n'

def test_pex_root():
  with temporary_dir() as tmp_home:
    with environment_as(HOME=tmp_home):
      with temporary_dir() as td:
        with temporary_dir() as output_dir:
          env = os.environ.copy()
          env['PEX_INTERPRETER'] = '1'

          output_path = os.path.join(output_dir, 'pex.pex')
          args = ['pex', '-o', output_path, '--not-zip-safe', '--pex-root={0}'.format(td)]
          results = run_pex_command(args=args, env=env)
          results.assert_success()
          assert ['pex.pex'] == os.listdir(output_dir), 'Expected built pex file.'
          assert [] == os.listdir(tmp_home), 'Expected empty temp home dir.'
          assert 'build' in os.listdir(td), 'Expected build directory in tmp pex root.'

 def _call(self, cmd, *args, **kwargs):
   """Runs hadoop fs command  with the given command and args.
   Checks the result of the call by default but this can be disabled with check=False.
   """
   cmd = ['hadoop', '--config', self._config, 'dfs', cmd] + list(args)
   heapsize = str(int(self._heap_limit.as_(Data.MB)))
   with environment_as(HADOOP_HEAPSIZE=heapsize):
     if kwargs.get('check'):
       return self._cmd_class.check_call(cmd)
     elif kwargs.get('return_output'):
       return self._cmd_class.execute_and_get_output(cmd)
     elif kwargs.get('supress_output'):
       return self._cmd_class.execute_suppress_stdout(cmd)
     else:
       return self._cmd_class.execute(cmd)

def write_and_run_simple_pex(inheriting=False):
  """Write a pex file that contains an executable entry point

  :param inheriting: whether this pex should inherit site-packages paths
  :type inheriting: bool
  """
  with temporary_dir() as td:
    pex_path = os.path.join(td, 'show_path.pex')
    with open(os.path.join(td, 'exe.py'), 'w') as fp:
      fp.write('')  # No contents, we just want the startup messages

    pb = PEXBuilder(path=td, preamble=None)
    pb.info.inherit_path = inheriting
    pb.set_executable(os.path.join(td, 'exe.py'))
    pb.freeze()
    pb.build(pex_path)
    with environment_as(PEX_VERBOSE='1'):
      yield run_simple_pex(pex_path)[0]