script_helper.py 5.2 KB
Newer Older
Nick Coghlan's avatar
Nick Coghlan committed
1 2 3 4 5
# Common utility functions used by various script execution tests
#  e.g. test_cmd_line, test_cmd_line_script and test_runpy

import sys
import os
6
import re
Nick Coghlan's avatar
Nick Coghlan committed
7 8 9 10 11 12 13 14
import os.path
import tempfile
import subprocess
import py_compile
import contextlib
import shutil
import zipfile

Barry Warsaw's avatar
Barry Warsaw committed
15
from imp import source_from_cache
16
from test.support import make_legacy_pyc, strip_python_stderr
Barry Warsaw's avatar
Barry Warsaw committed
17

Nick Coghlan's avatar
Nick Coghlan committed
18
# Executing the interpreter in a subprocess
19 20
def _assert_python(expected_success, *args, **env_vars):
    cmd_line = [sys.executable]
21
    if not env_vars:
22
        cmd_line.append('-E')
Nick Coghlan's avatar
Nick Coghlan committed
23
    cmd_line.extend(args)
24 25 26 27
    # Need to preserve the original environment, for in-place testing of
    # shared library builds.
    env = os.environ.copy()
    env.update(env_vars)
28
    p = subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
29 30
                         stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                         env=env)
31 32 33 34
    try:
        out, err = p.communicate()
    finally:
        subprocess._cleanup()
35 36
        p.stdout.close()
        p.stderr.close()
37
    rc = p.returncode
38
    err =  strip_python_stderr(err)
39 40 41 42 43 44
    if (rc and expected_success) or (not rc and not expected_success):
        raise AssertionError(
            "Process return code is %d, "
            "stderr follows:\n%s" % (rc, err.decode('ascii', 'ignore')))
    return rc, out, err

45 46 47 48 49 50
def assert_python_ok(*args, **env_vars):
    """
    Assert that running the interpreter with `args` and optional environment
    variables `env_vars` is ok and return a (return code, stdout, stderr) tuple.
    """
    return _assert_python(True, *args, **env_vars)
51

52 53 54 55 56 57
def assert_python_failure(*args, **env_vars):
    """
    Assert that running the interpreter with `args` and optional environment
    variables `env_vars` fails and return a (return code, stdout, stderr) tuple.
    """
    return _assert_python(False, *args, **env_vars)
Nick Coghlan's avatar
Nick Coghlan committed
58

59
def spawn_python(*args, **kw):
Nick Coghlan's avatar
Nick Coghlan committed
60 61 62
    cmd_line = [sys.executable, '-E']
    cmd_line.extend(args)
    return subprocess.Popen(cmd_line, stdin=subprocess.PIPE,
63 64
                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            **kw)
Nick Coghlan's avatar
Nick Coghlan committed
65 66 67 68 69 70

def kill_python(p):
    p.stdin.close()
    data = p.stdout.read()
    p.stdout.close()
    # try to cleanup the child so we don't appear to leak when running
71 72
    # with regrtest -R.
    p.wait()
Nick Coghlan's avatar
Nick Coghlan committed
73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88
    subprocess._cleanup()
    return data

# Script creation utilities
@contextlib.contextmanager
def temp_dir():
    dirname = tempfile.mkdtemp()
    dirname = os.path.realpath(dirname)
    try:
        yield dirname
    finally:
        shutil.rmtree(dirname)

def make_script(script_dir, script_basename, source):
    script_filename = script_basename+os.extsep+'py'
    script_name = os.path.join(script_dir, script_filename)
89 90
    # The script should be encoded to UTF-8, the default string encoding
    script_file = open(script_name, 'w', encoding='utf-8')
Nick Coghlan's avatar
Nick Coghlan committed
91 92 93 94 95 96 97 98 99
    script_file.write(source)
    script_file.close()
    return script_name

def make_zip_script(zip_dir, zip_basename, script_name, name_in_zip=None):
    zip_filename = zip_basename+os.extsep+'zip'
    zip_name = os.path.join(zip_dir, zip_filename)
    zip_file = zipfile.ZipFile(zip_name, 'w')
    if name_in_zip is None:
Barry Warsaw's avatar
Barry Warsaw committed
100 101 102 103 104 105 106
        parts = script_name.split(os.sep)
        if len(parts) >= 2 and parts[-2] == '__pycache__':
            legacy_pyc = make_legacy_pyc(source_from_cache(script_name))
            name_in_zip = os.path.basename(legacy_pyc)
            script_name = legacy_pyc
        else:
            name_in_zip = os.path.basename(script_name)
Nick Coghlan's avatar
Nick Coghlan committed
107 108
    zip_file.write(script_name, name_in_zip)
    zip_file.close()
Florent Xicluna's avatar
Florent Xicluna committed
109
    #if test.support.verbose:
Nick Coghlan's avatar
Nick Coghlan committed
110 111 112 113 114 115
    #    zip_file = zipfile.ZipFile(zip_name, 'r')
    #    print 'Contents of %r:' % zip_name
    #    zip_file.printdir()
    #    zip_file.close()
    return zip_name, os.path.join(zip_name, name_in_zip)

116
def make_pkg(pkg_dir, init_source=''):
Nick Coghlan's avatar
Nick Coghlan committed
117
    os.mkdir(pkg_dir)
118
    make_script(pkg_dir, '__init__', init_source)
Nick Coghlan's avatar
Nick Coghlan committed
119 120 121 122 123 124 125 126 127 128

def make_zip_pkg(zip_dir, zip_basename, pkg_name, script_basename,
                 source, depth=1, compiled=False):
    unlink = []
    init_name = make_script(zip_dir, '__init__', '')
    unlink.append(init_name)
    init_basename = os.path.basename(init_name)
    script_name = make_script(zip_dir, script_basename, source)
    unlink.append(script_name)
    if compiled:
Barry Warsaw's avatar
Barry Warsaw committed
129 130
        init_name = py_compile(init_name, doraise=True)
        script_name = py_compile(script_name, doraise=True)
Nick Coghlan's avatar
Nick Coghlan committed
131 132 133 134 135 136 137 138 139 140 141 142 143
        unlink.extend((init_name, script_name))
    pkg_names = [os.sep.join([pkg_name]*i) for i in range(1, depth+1)]
    script_name_in_zip = os.path.join(pkg_names[-1], os.path.basename(script_name))
    zip_filename = zip_basename+os.extsep+'zip'
    zip_name = os.path.join(zip_dir, zip_filename)
    zip_file = zipfile.ZipFile(zip_name, 'w')
    for name in pkg_names:
        init_name_in_zip = os.path.join(name, init_basename)
        zip_file.write(init_name, init_name_in_zip)
    zip_file.write(script_name, script_name_in_zip)
    zip_file.close()
    for name in unlink:
        os.unlink(name)
Florent Xicluna's avatar
Florent Xicluna committed
144
    #if test.support.verbose:
Nick Coghlan's avatar
Nick Coghlan committed
145 146 147 148 149
    #    zip_file = zipfile.ZipFile(zip_name, 'r')
    #    print 'Contents of %r:' % zip_name
    #    zip_file.printdir()
    #    zip_file.close()
    return zip_name, os.path.join(zip_name, script_name_in_zip)