script_helper.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170
  1. # Common utility functions used by various script execution tests
  2. # e.g. test_cmd_line, test_cmd_line_script and test_runpy
  3. import sys
  4. import os
  5. import re
  6. import os.path
  7. import tempfile
  8. import subprocess
  9. import py_compile
  10. import contextlib
  11. import shutil
  12. try:
  13. import zipfile
  14. except ImportError:
  15. # If Python is build without Unicode support, importing _io will
  16. # fail, which, in turn, means that zipfile cannot be imported
  17. # Most of this module can then still be used.
  18. pass
  19. from test.test_support import strip_python_stderr
  20. # Executing the interpreter in a subprocess
  21. def _assert_python(expected_success, *args, **env_vars):
  22. cmd_line = [sys.executable]
  23. if not env_vars:
  24. cmd_line.append('-E')
  25. cmd_line.extend(args)
  26. # Need to preserve the original environment, for in-place testing of
  27. # shared library builds.
  28. env = os.environ.copy()
  29. env.update(env_vars)
  30. p = subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
  31. stdout=subprocess.PIPE, stderr=subprocess.PIPE,
  32. env=env)
  33. try:
  34. out, err = p.communicate()
  35. finally:
  36. subprocess._cleanup()
  37. p.stdout.close()
  38. p.stderr.close()
  39. rc = p.returncode
  40. err = strip_python_stderr(err)
  41. if (rc and expected_success) or (not rc and not expected_success):
  42. raise AssertionError(
  43. "Process return code is %d, "
  44. "stderr follows:\n%s" % (rc, err.decode('ascii', 'ignore')))
  45. return rc, out, err
  46. def assert_python_ok(*args, **env_vars):
  47. """
  48. Assert that running the interpreter with `args` and optional environment
  49. variables `env_vars` is ok and return a (return code, stdout, stderr) tuple.
  50. """
  51. return _assert_python(True, *args, **env_vars)
  52. def assert_python_failure(*args, **env_vars):
  53. """
  54. Assert that running the interpreter with `args` and optional environment
  55. variables `env_vars` fails and return a (return code, stdout, stderr) tuple.
  56. """
  57. return _assert_python(False, *args, **env_vars)
  58. def python_exit_code(*args):
  59. cmd_line = [sys.executable, '-E']
  60. cmd_line.extend(args)
  61. with open(os.devnull, 'w') as devnull:
  62. return subprocess.call(cmd_line, stdout=devnull,
  63. stderr=subprocess.STDOUT)
  64. def spawn_python(*args, **kwargs):
  65. cmd_line = [sys.executable, '-E']
  66. cmd_line.extend(args)
  67. return subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
  68. stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
  69. **kwargs)
  70. def kill_python(p):
  71. p.stdin.close()
  72. data = p.stdout.read()
  73. p.stdout.close()
  74. # try to cleanup the child so we don't appear to leak when running
  75. # with regrtest -R.
  76. p.wait()
  77. subprocess._cleanup()
  78. return data
  79. def run_python(*args, **kwargs):
  80. if __debug__:
  81. p = spawn_python(*args, **kwargs)
  82. else:
  83. p = spawn_python('-O', *args, **kwargs)
  84. stdout_data = kill_python(p)
  85. return p.wait(), stdout_data
  86. # Script creation utilities
  87. @contextlib.contextmanager
  88. def temp_dir():
  89. dirname = tempfile.mkdtemp()
  90. dirname = os.path.realpath(dirname)
  91. try:
  92. yield dirname
  93. finally:
  94. shutil.rmtree(dirname)
  95. def make_script(script_dir, script_basename, source):
  96. script_filename = script_basename+os.extsep+'py'
  97. script_name = os.path.join(script_dir, script_filename)
  98. script_file = open(script_name, 'w')
  99. script_file.write(source)
  100. script_file.close()
  101. return script_name
  102. def compile_script(script_name):
  103. py_compile.compile(script_name, doraise=True)
  104. if __debug__:
  105. compiled_name = script_name + 'c'
  106. else:
  107. compiled_name = script_name + 'o'
  108. return compiled_name
  109. def make_zip_script(zip_dir, zip_basename, script_name, name_in_zip=None):
  110. zip_filename = zip_basename+os.extsep+'zip'
  111. zip_name = os.path.join(zip_dir, zip_filename)
  112. zip_file = zipfile.ZipFile(zip_name, 'w')
  113. if name_in_zip is None:
  114. name_in_zip = os.path.basename(script_name)
  115. zip_file.write(script_name, name_in_zip)
  116. zip_file.close()
  117. #if test.test_support.verbose:
  118. # zip_file = zipfile.ZipFile(zip_name, 'r')
  119. # print 'Contents of %r:' % zip_name
  120. # zip_file.printdir()
  121. # zip_file.close()
  122. return zip_name, os.path.join(zip_name, name_in_zip)
  123. def make_pkg(pkg_dir, init_source=''):
  124. os.mkdir(pkg_dir)
  125. make_script(pkg_dir, '__init__', init_source)
  126. def make_zip_pkg(zip_dir, zip_basename, pkg_name, script_basename,
  127. source, depth=1, compiled=False):
  128. unlink = []
  129. init_name = make_script(zip_dir, '__init__', '')
  130. unlink.append(init_name)
  131. init_basename = os.path.basename(init_name)
  132. script_name = make_script(zip_dir, script_basename, source)
  133. unlink.append(script_name)
  134. if compiled:
  135. init_name = compile_script(init_name)
  136. script_name = compile_script(script_name)
  137. unlink.extend((init_name, script_name))
  138. pkg_names = [os.sep.join([pkg_name]*i) for i in range(1, depth+1)]
  139. script_name_in_zip = os.path.join(pkg_names[-1], os.path.basename(script_name))
  140. zip_filename = zip_basename+os.extsep+'zip'
  141. zip_name = os.path.join(zip_dir, zip_filename)
  142. zip_file = zipfile.ZipFile(zip_name, 'w')
  143. for name in pkg_names:
  144. init_name_in_zip = os.path.join(name, init_basename)
  145. zip_file.write(init_name, init_name_in_zip)
  146. zip_file.write(script_name, script_name_in_zip)
  147. zip_file.close()
  148. for name in unlink:
  149. os.unlink(name)
  150. #if test.test_support.verbose:
  151. # zip_file = zipfile.ZipFile(zip_name, 'r')
  152. # print 'Contents of %r:' % zip_name
  153. # zip_file.printdir()
  154. # zip_file.close()
  155. return zip_name, os.path.join(zip_name, script_name_in_zip)