blob: 2bdcd9681e65a3dfea00e6a10d791eff8384b059 [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 implementation.
Reused across infra/ and infra_internal/
import os
import signal
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
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)."""
if not is_in_venv(env_path):
if RUN_PY_RECURSION_BLOCKER in os.environ:
sys.stderr.write('TOO MUCH RECURSION IN RUN.PY\n')
# not in the venv
if sys.platform.startswith('win'):
python = os.path.join(env_path, 'Scripts', 'python.exe')
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)
args = [python, script] + sys.argv[1:]
if sys.platform == 'win32':
# On Windows, os.execv spawns a child process and exits immediately with
# status zero, without waiting for it to finish. This confuses the
# recipe engine, which loses the stdout of the process and reports that
# the step has finished successfully as soon as the child process
# spawns, regardless of whether the child process fails.
# Our alternative implementation below waits on the child process, and
# exits with its return status, solving both of the above problems. It
# also works better when running from the console, as you don't get an
# immediate return back to the command prompt along with interleaved
# output from the child.
# This is a well-known Windows Python issue going back to at least 2001.
# See for more details.
signal.signal(signal.SIGBREAK, signal.SIG_IGN)
signal.signal(signal.SIGINT, signal.SIG_IGN)
signal.signal(signal.SIGTERM, signal.SIG_IGN)
os._exit(os.spawnv(os.P_WAIT, python, args))
os.execv(python, args)
sys.stderr.write('Exec is busted :(\n')
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.')
# In case some poor script ends up calling, 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 '' 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:
usage: %s.<> [args for tool]
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):
return orig_parse_args(*args, **kwargs)
argparse.ArgumentParser.parse_args = new_parse_args
# remove the module from sys.argv
del sys.argv[1]
runpy.run_module(args[0], run_name='__main__', alter_sys=True)