Lib/test/script_helper.py

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