1# Test the runpy module
2import unittest
3import os
4import os.path
5import sys
6import re
7import tempfile
8from test.test_support import verbose, run_unittest, forget
9from test.script_helper import (temp_dir, make_script, compile_script,
10                                make_pkg, make_zip_script, make_zip_pkg)
11
12
13from runpy import _run_code, _run_module_code, run_module, run_path
14# Note: This module can't safely test _run_module_as_main as it
15# runs its tests in the current process, which would mess with the
16# real __main__ module (usually test.regrtest)
17# See test_cmd_line_script for a test that executes that code path
18
19# Set up the test code and expected results
20
21class RunModuleCodeTest(unittest.TestCase):
22    """Unit tests for runpy._run_code and runpy._run_module_code"""
23
24    expected_result = ["Top level assignment", "Lower level reference"]
25    test_source = (
26        "# Check basic code execution\n"
27        "result = ['Top level assignment']\n"
28        "def f():\n"
29        "    result.append('Lower level reference')\n"
30        "f()\n"
31        "# Check the sys module\n"
32        "import sys\n"
33        "run_argv0 = sys.argv[0]\n"
34        "run_name_in_sys_modules = __name__ in sys.modules\n"
35        "if run_name_in_sys_modules:\n"
36        "   module_in_sys_modules = globals() is sys.modules[__name__].__dict__\n"
37        "# Check nested operation\n"
38        "import runpy\n"
39        "nested = runpy._run_module_code('x=1\\n', mod_name='<run>')\n"
40    )
41
42    def test_run_code(self):
43        saved_argv0 = sys.argv[0]
44        d = _run_code(self.test_source, {})
45        self.assertEqual(d["result"], self.expected_result)
46        self.assertIs(d["__name__"], None)
47        self.assertIs(d["__file__"], None)
48        self.assertIs(d["__loader__"], None)
49        self.assertIs(d["__package__"], None)
50        self.assertIs(d["run_argv0"], saved_argv0)
51        self.assertNotIn("run_name", d)
52        self.assertIs(sys.argv[0], saved_argv0)
53
54    def test_run_module_code(self):
55        initial = object()
56        name = "<Nonsense>"
57        file = "Some other nonsense"
58        loader = "Now you're just being silly"
59        package = '' # Treat as a top level module
60        d1 = dict(initial=initial)
61        saved_argv0 = sys.argv[0]
62        d2 = _run_module_code(self.test_source,
63                              d1,
64                              name,
65                              file,
66                              loader,
67                              package)
68        self.assertNotIn("result", d1)
69        self.assertIs(d2["initial"], initial)
70        self.assertEqual(d2["result"], self.expected_result)
71        self.assertEqual(d2["nested"]["x"], 1)
72        self.assertIs(d2["__name__"], name)
73        self.assertTrue(d2["run_name_in_sys_modules"])
74        self.assertTrue(d2["module_in_sys_modules"])
75        self.assertIs(d2["__file__"], file)
76        self.assertIs(d2["run_argv0"], file)
77        self.assertIs(d2["__loader__"], loader)
78        self.assertIs(d2["__package__"], package)
79        self.assertIs(sys.argv[0], saved_argv0)
80        self.assertNotIn(name, sys.modules)
81
82
83class RunModuleTest(unittest.TestCase):
84    """Unit tests for runpy.run_module"""
85
86    def expect_import_error(self, mod_name):
87        try:
88            run_module(mod_name)
89        except ImportError:
90            pass
91        else:
92            self.fail("Expected import error for " + mod_name)
93
94    def test_invalid_names(self):
95        # Builtin module
96        self.expect_import_error("sys")
97        # Non-existent modules
98        self.expect_import_error("sys.imp.eric")
99        self.expect_import_error("os.path.half")
100        self.expect_import_error("a.bee")
101        self.expect_import_error(".howard")
102        self.expect_import_error("..eaten")
103        # Package without __main__.py
104        self.expect_import_error("multiprocessing")
105
106    def test_library_module(self):
107        run_module("runpy")
108
109    def _add_pkg_dir(self, pkg_dir):
110        os.mkdir(pkg_dir)
111        pkg_fname = os.path.join(pkg_dir, "__init__"+os.extsep+"py")
112        pkg_file = open(pkg_fname, "w")
113        pkg_file.close()
114        return pkg_fname
115
116    def _make_pkg(self, source, depth, mod_base="runpy_test"):
117        pkg_name = "__runpy_pkg__"
118        test_fname = mod_base+os.extsep+"py"
119        pkg_dir = sub_dir = tempfile.mkdtemp()
120        if verbose: print "  Package tree in:", sub_dir
121        sys.path.insert(0, pkg_dir)
122        if verbose: print "  Updated sys.path:", sys.path[0]
123        for i in range(depth):
124            sub_dir = os.path.join(sub_dir, pkg_name)
125            pkg_fname = self._add_pkg_dir(sub_dir)
126            if verbose: print "  Next level in:", sub_dir
127            if verbose: print "  Created:", pkg_fname
128        mod_fname = os.path.join(sub_dir, test_fname)
129        mod_file = open(mod_fname, "w")
130        mod_file.write(source)
131        mod_file.close()
132        if verbose: print "  Created:", mod_fname
133        mod_name = (pkg_name+".")*depth + mod_base
134        return pkg_dir, mod_fname, mod_name
135
136    def _del_pkg(self, top, depth, mod_name):
137        for entry in list(sys.modules):
138            if entry.startswith("__runpy_pkg__"):
139                del sys.modules[entry]
140        if verbose: print "  Removed sys.modules entries"
141        del sys.path[0]
142        if verbose: print "  Removed sys.path entry"
143        for root, dirs, files in os.walk(top, topdown=False):
144            for name in files:
145                try:
146                    os.remove(os.path.join(root, name))
147                except OSError, ex:
148                    if verbose: print ex # Persist with cleaning up
149            for name in dirs:
150                fullname = os.path.join(root, name)
151                try:
152                    os.rmdir(fullname)
153                except OSError, ex:
154                    if verbose: print ex # Persist with cleaning up
155        try:
156            os.rmdir(top)
157            if verbose: print "  Removed package tree"
158        except OSError, ex:
159            if verbose: print ex # Persist with cleaning up
160
161    def _check_module(self, depth):
162        pkg_dir, mod_fname, mod_name = (
163               self._make_pkg("x=1\n", depth))
164        forget(mod_name)
165        try:
166            if verbose: print "Running from source:", mod_name
167            d1 = run_module(mod_name) # Read from source
168            self.assertIn("x", d1)
169            self.assertTrue(d1["x"] == 1)
170            del d1 # Ensure __loader__ entry doesn't keep file open
171            __import__(mod_name)
172            os.remove(mod_fname)
173            if not sys.dont_write_bytecode:
174                if verbose: print "Running from compiled:", mod_name
175                d2 = run_module(mod_name) # Read from bytecode
176                self.assertIn("x", d2)
177                self.assertTrue(d2["x"] == 1)
178                del d2 # Ensure __loader__ entry doesn't keep file open
179        finally:
180            self._del_pkg(pkg_dir, depth, mod_name)
181        if verbose: print "Module executed successfully"
182
183    def _check_package(self, depth):
184        pkg_dir, mod_fname, mod_name = (
185               self._make_pkg("x=1\n", depth, "__main__"))
186        pkg_name, _, _ = mod_name.rpartition(".")
187        forget(mod_name)
188        try:
189            if verbose: print "Running from source:", pkg_name
190            d1 = run_module(pkg_name) # Read from source
191            self.assertIn("x", d1)
192            self.assertTrue(d1["x"] == 1)
193            del d1 # Ensure __loader__ entry doesn't keep file open
194            __import__(mod_name)
195            os.remove(mod_fname)
196            if not sys.dont_write_bytecode:
197                if verbose: print "Running from compiled:", pkg_name
198                d2 = run_module(pkg_name) # Read from bytecode
199                self.assertIn("x", d2)
200                self.assertTrue(d2["x"] == 1)
201                del d2 # Ensure __loader__ entry doesn't keep file open
202        finally:
203            self._del_pkg(pkg_dir, depth, pkg_name)
204        if verbose: print "Package executed successfully"
205
206    def _add_relative_modules(self, base_dir, source, depth):
207        if depth <= 1:
208            raise ValueError("Relative module test needs depth > 1")
209        pkg_name = "__runpy_pkg__"
210        module_dir = base_dir
211        for i in range(depth):
212            parent_dir = module_dir
213            module_dir = os.path.join(module_dir, pkg_name)
214        # Add sibling module
215        sibling_fname = os.path.join(module_dir, "sibling"+os.extsep+"py")
216        sibling_file = open(sibling_fname, "w")
217        sibling_file.close()
218        if verbose: print "  Added sibling module:", sibling_fname
219        # Add nephew module
220        uncle_dir = os.path.join(parent_dir, "uncle")
221        self._add_pkg_dir(uncle_dir)
222        if verbose: print "  Added uncle package:", uncle_dir
223        cousin_dir = os.path.join(uncle_dir, "cousin")
224        self._add_pkg_dir(cousin_dir)
225        if verbose: print "  Added cousin package:", cousin_dir
226        nephew_fname = os.path.join(cousin_dir, "nephew"+os.extsep+"py")
227        nephew_file = open(nephew_fname, "w")
228        nephew_file.close()
229        if verbose: print "  Added nephew module:", nephew_fname
230
231    def _check_relative_imports(self, depth, run_name=None):
232        contents = r"""\
233from __future__ import absolute_import
234from . import sibling
235from ..uncle.cousin import nephew
236"""
237        pkg_dir, mod_fname, mod_name = (
238               self._make_pkg(contents, depth))
239        try:
240            self._add_relative_modules(pkg_dir, contents, depth)
241            pkg_name = mod_name.rpartition('.')[0]
242            if verbose: print "Running from source:", mod_name
243            d1 = run_module(mod_name, run_name=run_name) # Read from source
244            self.assertIn("__package__", d1)
245            self.assertTrue(d1["__package__"] == pkg_name)
246            self.assertIn("sibling", d1)
247            self.assertIn("nephew", d1)
248            del d1 # Ensure __loader__ entry doesn't keep file open
249            __import__(mod_name)
250            os.remove(mod_fname)
251            if not sys.dont_write_bytecode:
252                if verbose: print "Running from compiled:", mod_name
253                d2 = run_module(mod_name, run_name=run_name) # Read from bytecode
254                self.assertIn("__package__", d2)
255                self.assertTrue(d2["__package__"] == pkg_name)
256                self.assertIn("sibling", d2)
257                self.assertIn("nephew", d2)
258                del d2 # Ensure __loader__ entry doesn't keep file open
259        finally:
260            self._del_pkg(pkg_dir, depth, mod_name)
261        if verbose: print "Module executed successfully"
262
263    def test_run_module(self):
264        for depth in range(4):
265            if verbose: print "Testing package depth:", depth
266            self._check_module(depth)
267
268    def test_run_package(self):
269        for depth in range(1, 4):
270            if verbose: print "Testing package depth:", depth
271            self._check_package(depth)
272
273    def test_explicit_relative_import(self):
274        for depth in range(2, 5):
275            if verbose: print "Testing relative imports at depth:", depth
276            self._check_relative_imports(depth)
277
278    def test_main_relative_import(self):
279        for depth in range(2, 5):
280            if verbose: print "Testing main relative imports at depth:", depth
281            self._check_relative_imports(depth, "__main__")
282
283
284class RunPathTest(unittest.TestCase):
285    """Unit tests for runpy.run_path"""
286    # Based on corresponding tests in test_cmd_line_script
287
288    test_source = """\
289# Script may be run with optimisation enabled, so don't rely on assert
290# statements being executed
291def assertEqual(lhs, rhs):
292    if lhs != rhs:
293        raise AssertionError('%r != %r' % (lhs, rhs))
294def assertIs(lhs, rhs):
295    if lhs is not rhs:
296        raise AssertionError('%r is not %r' % (lhs, rhs))
297# Check basic code execution
298result = ['Top level assignment']
299def f():
300    result.append('Lower level reference')
301f()
302assertEqual(result, ['Top level assignment', 'Lower level reference'])
303# Check the sys module
304import sys
305assertIs(globals(), sys.modules[__name__].__dict__)
306argv0 = sys.argv[0]
307"""
308
309    def _make_test_script(self, script_dir, script_basename, source=None):
310        if source is None:
311            source = self.test_source
312        return make_script(script_dir, script_basename, source)
313
314    def _check_script(self, script_name, expected_name, expected_file,
315                            expected_argv0, expected_package):
316        result = run_path(script_name)
317        self.assertEqual(result["__name__"], expected_name)
318        self.assertEqual(result["__file__"], expected_file)
319        self.assertIn("argv0", result)
320        self.assertEqual(result["argv0"], expected_argv0)
321        self.assertEqual(result["__package__"], expected_package)
322
323    def _check_import_error(self, script_name, msg):
324        msg = re.escape(msg)
325        self.assertRaisesRegexp(ImportError, msg, run_path, script_name)
326
327    def test_basic_script(self):
328        with temp_dir() as script_dir:
329            mod_name = 'script'
330            script_name = self._make_test_script(script_dir, mod_name)
331            self._check_script(script_name, "<run_path>", script_name,
332                               script_name, None)
333
334    def test_script_compiled(self):
335        with temp_dir() as script_dir:
336            mod_name = 'script'
337            script_name = self._make_test_script(script_dir, mod_name)
338            compiled_name = compile_script(script_name)
339            os.remove(script_name)
340            self._check_script(compiled_name, "<run_path>", compiled_name,
341                               compiled_name, None)
342
343    def test_directory(self):
344        with temp_dir() as script_dir:
345            mod_name = '__main__'
346            script_name = self._make_test_script(script_dir, mod_name)
347            self._check_script(script_dir, "<run_path>", script_name,
348                               script_dir, '')
349
350    def test_directory_compiled(self):
351        with temp_dir() as script_dir:
352            mod_name = '__main__'
353            script_name = self._make_test_script(script_dir, mod_name)
354            compiled_name = compile_script(script_name)
355            os.remove(script_name)
356            self._check_script(script_dir, "<run_path>", compiled_name,
357                               script_dir, '')
358
359    def test_directory_error(self):
360        with temp_dir() as script_dir:
361            mod_name = 'not_main'
362            script_name = self._make_test_script(script_dir, mod_name)
363            msg = "can't find '__main__' module in %r" % script_dir
364            self._check_import_error(script_dir, msg)
365
366    def test_zipfile(self):
367        with temp_dir() as script_dir:
368            mod_name = '__main__'
369            script_name = self._make_test_script(script_dir, mod_name)
370            zip_name, fname = make_zip_script(script_dir, 'test_zip', script_name)
371            self._check_script(zip_name, "<run_path>", fname, zip_name, '')
372
373    def test_zipfile_compiled(self):
374        with temp_dir() as script_dir:
375            mod_name = '__main__'
376            script_name = self._make_test_script(script_dir, mod_name)
377            compiled_name = compile_script(script_name)
378            zip_name, fname = make_zip_script(script_dir, 'test_zip', compiled_name)
379            self._check_script(zip_name, "<run_path>", fname, zip_name, '')
380
381    def test_zipfile_error(self):
382        with temp_dir() as script_dir:
383            mod_name = 'not_main'
384            script_name = self._make_test_script(script_dir, mod_name)
385            zip_name, fname = make_zip_script(script_dir, 'test_zip', script_name)
386            msg = "can't find '__main__' module in %r" % zip_name
387            self._check_import_error(zip_name, msg)
388
389    def test_main_recursion_error(self):
390        with temp_dir() as script_dir, temp_dir() as dummy_dir:
391            mod_name = '__main__'
392            source = ("import runpy\n"
393                      "runpy.run_path(%r)\n") % dummy_dir
394            script_name = self._make_test_script(script_dir, mod_name, source)
395            zip_name, fname = make_zip_script(script_dir, 'test_zip', script_name)
396            msg = "recursion depth exceeded"
397            self.assertRaisesRegexp(RuntimeError, msg, run_path, zip_name)
398
399
400
401def test_main():
402    run_unittest(RunModuleCodeTest, RunModuleTest, RunPathTest)
403
404if __name__ == "__main__":
405    test_main()
406