generate_wrapper.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216
  1. #!/usr/bin/env python3
  2. # Copyright 2019 The Chromium Authors. All rights reserved.
  3. # Use of this source code is governed by a BSD-style license that can be
  4. # found in the LICENSE file.
  5. """Wraps an executable and any provided arguments into an executable script."""
  6. import argparse
  7. import os
  8. import sys
  9. import textwrap
  10. # The bash template passes the python script into vpython via stdin.
  11. # The interpreter doesn't know about the script, so we have bash
  12. # inject the script location.
  13. BASH_TEMPLATE = textwrap.dedent("""\
  14. #!/usr/bin/env vpython3
  15. _SCRIPT_LOCATION = __file__
  16. {script}
  17. """)
  18. # The batch template reruns the batch script with vpython, with the -x
  19. # flag instructing the interpreter to ignore the first line. The interpreter
  20. # knows about the (batch) script in this case, so it can get the file location
  21. # directly.
  22. BATCH_TEMPLATE = textwrap.dedent("""\
  23. @SETLOCAL ENABLEDELAYEDEXPANSION \
  24. & vpython3.bat -x "%~f0" %* \
  25. & EXIT /B !ERRORLEVEL!
  26. _SCRIPT_LOCATION = __file__
  27. {script}
  28. """)
  29. SCRIPT_TEMPLATES = {
  30. 'bash': BASH_TEMPLATE,
  31. 'batch': BATCH_TEMPLATE,
  32. }
  33. PY_TEMPLATE = textwrap.dedent("""\
  34. import os
  35. import re
  36. import shlex
  37. import subprocess
  38. import sys
  39. _WRAPPED_PATH_RE = re.compile(r'@WrappedPath\(([^)]+)\)')
  40. _PATH_TO_OUTPUT_DIR = '{path_to_output_dir}'
  41. _SCRIPT_DIR = os.path.dirname(os.path.realpath(_SCRIPT_LOCATION))
  42. def ExpandWrappedPath(arg):
  43. m = _WRAPPED_PATH_RE.match(arg)
  44. if m:
  45. relpath = os.path.join(
  46. os.path.relpath(_SCRIPT_DIR), _PATH_TO_OUTPUT_DIR, m.group(1))
  47. npath = os.path.normpath(relpath)
  48. if os.path.sep not in npath:
  49. # If the original path points to something in the current directory,
  50. # returning the normalized version of it can be a problem.
  51. # normpath() strips off the './' part of the path
  52. # ('./foo' becomes 'foo'), which can be a problem if the result
  53. # is passed to something like os.execvp(); in that case
  54. # osexecvp() will search $PATH for the executable, rather than
  55. # just execing the arg directly, and if '.' isn't in $PATH, this
  56. # results in an error.
  57. #
  58. # So, we need to explicitly return './foo' (or '.\\foo' on windows)
  59. # instead of 'foo'.
  60. #
  61. # Hopefully there are no cases where this causes a problem; if
  62. # there are, we will either need to change the interface to
  63. # WrappedPath() somehow to distinguish between the two, or
  64. # somehow ensure that the wrapped executable doesn't hit cases
  65. # like this.
  66. return '.' + os.path.sep + npath
  67. return npath
  68. return arg
  69. def ExpandWrappedPaths(args):
  70. for i, arg in enumerate(args):
  71. args[i] = ExpandWrappedPath(arg)
  72. return args
  73. def FindIsolatedOutdir(raw_args):
  74. outdir = None
  75. i = 0
  76. remaining_args = []
  77. while i < len(raw_args):
  78. if raw_args[i] == '--isolated-outdir' and i < len(raw_args)-1:
  79. outdir = raw_args[i+1]
  80. i += 2
  81. elif raw_args[i].startswith('--isolated-outdir='):
  82. outdir = raw_args[i][len('--isolated-outdir='):]
  83. i += 1
  84. else:
  85. remaining_args.append(raw_args[i])
  86. i += 1
  87. if not outdir and 'ISOLATED_OUTDIR' in os.environ:
  88. outdir = os.environ['ISOLATED_OUTDIR']
  89. return outdir, remaining_args
  90. def InsertWrapperScriptArgs(args):
  91. if '--wrapper-script-args' in args:
  92. idx = args.index('--wrapper-script-args')
  93. args.insert(idx + 1, shlex.join(sys.argv))
  94. def FilterIsolatedOutdirBasedArgs(outdir, args):
  95. rargs = []
  96. i = 0
  97. while i < len(args):
  98. if 'ISOLATED_OUTDIR' in args[i]:
  99. if outdir:
  100. # Rewrite the arg.
  101. rargs.append(args[i].replace('${{ISOLATED_OUTDIR}}',
  102. outdir).replace(
  103. '$ISOLATED_OUTDIR', outdir))
  104. i += 1
  105. else:
  106. # Simply drop the arg.
  107. i += 1
  108. elif (not outdir and
  109. args[i].startswith('-') and
  110. '=' not in args[i] and
  111. i < len(args) - 1 and
  112. 'ISOLATED_OUTDIR' in args[i+1]):
  113. # Parsing this case is ambiguous; if we're given
  114. # `--foo $ISOLATED_OUTDIR` we can't tell if $ISOLATED_OUTDIR
  115. # is meant to be the value of foo, or if foo takes no argument
  116. # and $ISOLATED_OUTDIR is the first positional arg.
  117. #
  118. # We assume the former will be much more common, and so we
  119. # need to drop --foo and $ISOLATED_OUTDIR.
  120. i += 2
  121. else:
  122. rargs.append(args[i])
  123. i += 1
  124. return rargs
  125. def main(raw_args):
  126. executable_path = ExpandWrappedPath('{executable_path}')
  127. outdir, remaining_args = FindIsolatedOutdir(raw_args)
  128. args = {executable_args}
  129. InsertWrapperScriptArgs(args)
  130. args = FilterIsolatedOutdirBasedArgs(outdir, args)
  131. executable_args = ExpandWrappedPaths(args)
  132. cmd = [executable_path] + executable_args + remaining_args
  133. if executable_path.endswith('.py'):
  134. cmd = [sys.executable] + cmd
  135. return subprocess.call(cmd)
  136. if __name__ == '__main__':
  137. sys.exit(main(sys.argv[1:]))
  138. """)
  139. def Wrap(args):
  140. """Writes a wrapped script according to the provided arguments.
  141. Arguments:
  142. args: an argparse.Namespace object containing command-line arguments
  143. as parsed by a parser returned by CreateArgumentParser.
  144. """
  145. path_to_output_dir = os.path.relpath(
  146. args.output_directory,
  147. os.path.dirname(args.wrapper_script))
  148. with open(args.wrapper_script, 'w') as wrapper_script:
  149. py_contents = PY_TEMPLATE.format(
  150. path_to_output_dir=path_to_output_dir,
  151. executable_path=str(args.executable),
  152. executable_args=str(args.executable_args))
  153. template = SCRIPT_TEMPLATES[args.script_language]
  154. wrapper_script.write(template.format(script=py_contents))
  155. os.chmod(args.wrapper_script, 0o750)
  156. return 0
  157. def CreateArgumentParser():
  158. """Creates an argparse.ArgumentParser instance."""
  159. parser = argparse.ArgumentParser()
  160. parser.add_argument(
  161. '--executable',
  162. help='Executable to wrap.')
  163. parser.add_argument(
  164. '--wrapper-script',
  165. help='Path to which the wrapper script will be written.')
  166. parser.add_argument(
  167. '--output-directory',
  168. help='Path to the output directory.')
  169. parser.add_argument(
  170. '--script-language',
  171. choices=SCRIPT_TEMPLATES.keys(),
  172. help='Language in which the wrapper script will be written.')
  173. parser.add_argument(
  174. 'executable_args', nargs='*',
  175. help='Arguments to wrap into the executable.')
  176. return parser
  177. def main(raw_args):
  178. parser = CreateArgumentParser()
  179. args = parser.parse_args(raw_args)
  180. return Wrap(args)
  181. if __name__ == '__main__':
  182. sys.exit(main(sys.argv[1:]))