1# Common utility functions used by various script execution tests
2#  e.g. test_cmd_line, test_cmd_line_script and test_runpy
3
4import sys
5import os
6import re
7import os.path
8import tempfile
9import subprocess
10import py_compile
11import contextlib
12import shutil
13import zipfile
14
15from test.test_support import strip_python_stderr
16
17# Executing the interpreter in a subprocess
18def _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
44def 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
51def 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
58def 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
65def 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
72def 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
82def 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
92def temp_dir():
93    dirname = tempfile.mkdtemp()
94    dirname = os.path.realpath(dirname)
95    try:
96        yield dirname
97    finally:
98        shutil.rmtree(dirname)
99
100def 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
108def 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
116def 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
131def make_pkg(pkg_dir):
132    os.mkdir(pkg_dir)
133    make_script(pkg_dir, '__init__', '')
134
135def 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)
165