bootstrap/run_helper.py

#!/usr/bin/env python
# Copyright 2014 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.

"""Code supporting run.py implementation.

Reused across infra/run.py and infra_internal/run.py.
"""

import os
import sys


def is_in_venv(env_path):
  """True if already running in virtual env."""
  abs_prefix = os.path.abspath(sys.prefix)
  abs_env_path = os.path.abspath(env_path)
  if abs_prefix == abs_env_path:
    return True
  # Ordinarily os.path.abspath(sys.prefix) == env_path is enough. But it doesn't
  # work when virtual env is deployed as CIPD package. CIPD uses symlinks to
  # stage files into installation root. When booting venv, something (python
  # binary itself?) resolves the symlink ENV/bin/python to the target, making
  # sys.prefix look like "<root>/.cipd/.../ENV". Note that "<root>/ENV" is not
  # a symlink itself, but "<root>/ENV/bin/python" is.
  if sys.platform == 'win32':
    # TODO(vadimsh): Make it work for Win32 too.
    return False
  try:
    return os.path.samefile(
        os.path.join(abs_prefix, 'bin', 'python'),
        os.path.join(abs_env_path, 'bin', 'python'))
  except OSError:
    return False


def boot_venv(script, env_path):
  """Reexecs the top-level script in a virtualenv (if necessary)."""
  RUN_PY_RECURSION_BLOCKER = 'RUN_PY_RECURSION'

  if not is_in_venv(env_path):
    if RUN_PY_RECURSION_BLOCKER in os.environ:
      print >> sys.stderr, 'TOO MUCH RECURSION IN RUN.PY'
      sys.exit(-1)

    # not in the venv
    if sys.platform.startswith('win'):
      python = os.path.join(env_path, 'Scripts', 'python.exe')
    else:
      python = os.path.join(env_path, 'bin', 'python')
    if os.path.exists(python):
      os.environ[RUN_PY_RECURSION_BLOCKER] = "1"
      os.environ.pop('PYTHONPATH', None)
      os.execv(python, [python, script] + sys.argv[1:])
      print >> sys.stderr, "Exec is busted :("
      sys.exit(-1)  # should never reach

    print 'You must use the virtualenv in ENV for scripts in the infra repo.'
    print 'Running `gclient runhooks` will create this environment for you.'
    sys.exit(1)

  # In case some poor script ends up calling run.py, don't explode them.
  os.environ.pop(RUN_PY_RECURSION_BLOCKER, None)


def run_py_main(args, runpy_path, env_path, package):
  boot_venv(runpy_path, env_path)

  import argparse
  import runpy
  import shlex
  import textwrap

  import argcomplete

  # Impersonate the argcomplete 'protocol'
  completing = os.getenv('_ARGCOMPLETE') == '1'
  if completing:
    assert not args
    line = os.getenv('COMP_LINE')
    args = shlex.split(line)[1:]
    if len(args) == 1 and not line.endswith(' '):
      args = []

  if not args or not args[0].startswith('%s.' % package):
    commands = []
    for root, _, files in os.walk(package):
      if '__main__.py' in files:
        commands.append(root.replace(os.path.sep, '.'))
    commands = sorted(commands)

    if completing:
      # Argcomplete is listening for strings on fd 8
      with os.fdopen(8, 'wb') as f:
        print >> f, '\n'.join(commands)
      return

    print textwrap.dedent("""\
    usage: run.py %s.<module.path.to.tool> [args for tool]

    %s

    Available tools are:""") % (
        package, sys.modules['__main__'].__doc__.strip())
    for command in commands:
      print '  *', command
    return 1

  if completing:
    to_nuke = ' ' + args[0]
    os.environ['COMP_LINE'] = os.environ['COMP_LINE'].replace(to_nuke, '', 1)
    os.environ['COMP_POINT'] = str(int(os.environ['COMP_POINT']) - len(to_nuke))
    orig_parse_args = argparse.ArgumentParser.parse_args
    def new_parse_args(self, *args, **kwargs):
      argcomplete.autocomplete(self)
      return orig_parse_args(*args, **kwargs)
    argparse.ArgumentParser.parse_args = new_parse_args
  else:
    # remove the module from sys.argv
    del sys.argv[1]

  runpy.run_module(args[0], run_name='__main__', alter_sys=True)