blob: a2d1e8648e0d147b834f7267fa6c14c03958d68d [file] [log] [blame]
#!/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
os.chdir(os.path.dirname(runpy_path))
# 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)