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 os.path
   7 import tempfile
   8 import subprocess
   9 import py_compile
  10 import contextlib
  11 import shutil
  12 import zipfile
  13
  14 # Executing the interpreter in a subprocess
  15 def python_exit_code(*args):
  16     cmd_line = [sys.executable, '-E']
  17     cmd_line.extend(args)
  18     with open(os.devnull, 'w') as devnull:
  19         return subprocess.call(cmd_line, stdout=devnull,
  20                                 stderr=subprocess.STDOUT)
  21
  22 def spawn_python(*args, **kwargs):
  23     cmd_line = [sys.executable, '-E']
  24     cmd_line.extend(args)
  25     return subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
  26                             stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
  27                             **kwargs)
  28
  29 def kill_python(p):
  30     p.stdin.close()
  31     data = p.stdout.read()
  32     p.stdout.close()
  33     # try to cleanup the child so we don't appear to leak when running
  34     # with regrtest -R.
  35     p.wait()
  36     subprocess._cleanup()
  37     return data
  38
  39 def run_python(*args, **kwargs):
  40     if __debug__:
  41         p = spawn_python(*args, **kwargs)
  42     else:
  43         p = spawn_python('-O', *args, **kwargs)
  44     stdout_data = kill_python(p)
  45     return p.wait(), stdout_data
  46
  47 # Script creation utilities
  48 @contextlib.contextmanager
  49 def temp_dir():
  50     dirname = tempfile.mkdtemp()
  51     dirname = os.path.realpath(dirname)
  52     try:
  53         yield dirname
  54     finally:
  55         shutil.rmtree(dirname)
  56
  57 def make_script(script_dir, script_basename, source):
  58     script_filename = script_basename+os.extsep+'py'
  59     script_name = os.path.join(script_dir, script_filename)
  60     script_file = open(script_name, 'w')
  61     script_file.write(source)
  62     script_file.close()
  63     return script_name
  64
  65 def compile_script(script_name):
  66     py_compile.compile(script_name, doraise=True)
  67     if __debug__:
  68         compiled_name = script_name + 'c'
  69     else:
  70         compiled_name = script_name + 'o'
  71     return compiled_name
  72
  73 def make_zip_script(zip_dir, zip_basename, script_name, name_in_zip=None):
  74     zip_filename = zip_basename+os.extsep+'zip'
  75     zip_name = os.path.join(zip_dir, zip_filename)
  76     zip_file = zipfile.ZipFile(zip_name, 'w')
  77     if name_in_zip is None:
  78         name_in_zip = os.path.basename(script_name)
  79     zip_file.write(script_name, name_in_zip)
  80     zip_file.close()
  81     #if test.test_support.verbose:
  82     #    zip_file = zipfile.ZipFile(zip_name, 'r')
  83     #    print 'Contents of %r:' % zip_name
  84     #    zip_file.printdir()
  85     #    zip_file.close()
  86     return zip_name, os.path.join(zip_name, name_in_zip)
  87
  88 def make_pkg(pkg_dir):
  89     os.mkdir(pkg_dir)
  90     make_script(pkg_dir, '__init__', '')
  91
  92 def make_zip_pkg(zip_dir, zip_basename, pkg_name, script_basename,
  93                  source, depth=1, compiled=False):
  94     unlink = []
  95     init_name = make_script(zip_dir, '__init__', '')
  96     unlink.append(init_name)
  97     init_basename = os.path.basename(init_name)
  98     script_name = make_script(zip_dir, script_basename, source)
  99     unlink.append(script_name)
 100     if compiled:
 101         init_name = compile_script(init_name)
 102         script_name = compile_script(script_name)
 103         unlink.extend((init_name, script_name))
 104     pkg_names = [os.sep.join([pkg_name]*i) for i in range(1, depth+1)]
 105     script_name_in_zip = os.path.join(pkg_names[-1], os.path.basename(script_name))
 106     zip_filename = zip_basename+os.extsep+'zip'
 107     zip_name = os.path.join(zip_dir, zip_filename)
 108     zip_file = zipfile.ZipFile(zip_name, 'w')
 109     for name in pkg_names:
 110         init_name_in_zip = os.path.join(name, init_basename)
 111         zip_file.write(init_name, init_name_in_zip)
 112     zip_file.write(script_name, script_name_in_zip)
 113     zip_file.close()
 114     for name in unlink:
 115         os.unlink(name)
 116     #if test.test_support.verbose:
 117     #    zip_file = zipfile.ZipFile(zip_name, 'r')
 118     #    print 'Contents of %r:' % zip_name
 119     #    zip_file.printdir()
 120     #    zip_file.close()
 121     return zip_name, os.path.join(zip_name, script_name_in_zip)