1import compileall
2import imp
3import os
4import py_compile
5import shutil
6import struct
7import tempfile
8from test import test_support
9import unittest
10
11
12class CompileallTests(unittest.TestCase):
13
14    def setUp(self):
15        self.directory = tempfile.mkdtemp()
16        self.source_path = os.path.join(self.directory, '_test.py')
17        self.bc_path = self.source_path + ('c' if __debug__ else 'o')
18        with open(self.source_path, 'w') as file:
19            file.write('x = 123\n')
20        self.source_path2 = os.path.join(self.directory, '_test2.py')
21        self.bc_path2 = self.source_path2 + ('c' if __debug__ else 'o')
22        shutil.copyfile(self.source_path, self.source_path2)
23
24    def tearDown(self):
25        shutil.rmtree(self.directory)
26
27    def data(self):
28        with open(self.bc_path, 'rb') as file:
29            data = file.read(8)
30        mtime = int(os.stat(self.source_path).st_mtime)
31        compare = struct.pack('<4sl', imp.get_magic(), mtime)
32        return data, compare
33
34    def recreation_check(self, metadata):
35        """Check that compileall recreates bytecode when the new metadata is
36        used."""
37        if not hasattr(os, 'stat'):
38            return
39        py_compile.compile(self.source_path)
40        self.assertEqual(*self.data())
41        with open(self.bc_path, 'rb') as file:
42            bc = file.read()[len(metadata):]
43        with open(self.bc_path, 'wb') as file:
44            file.write(metadata)
45            file.write(bc)
46        self.assertNotEqual(*self.data())
47        compileall.compile_dir(self.directory, force=False, quiet=True)
48        self.assertTrue(*self.data())
49
50    def test_mtime(self):
51        # Test a change in mtime leads to a new .pyc.
52        self.recreation_check(struct.pack('<4sl', imp.get_magic(), 1))
53
54    def test_magic_number(self):
55        # Test a change in mtime leads to a new .pyc.
56        self.recreation_check(b'\0\0\0\0')
57
58    def test_compile_files(self):
59        # Test compiling a single file, and complete directory
60        for fn in (self.bc_path, self.bc_path2):
61            try:
62                os.unlink(fn)
63            except:
64                pass
65        compileall.compile_file(self.source_path, force=False, quiet=True)
66        self.assertTrue(os.path.isfile(self.bc_path) \
67                        and not os.path.isfile(self.bc_path2))
68        os.unlink(self.bc_path)
69        compileall.compile_dir(self.directory, force=False, quiet=True)
70        self.assertTrue(os.path.isfile(self.bc_path) \
71                        and os.path.isfile(self.bc_path2))
72        os.unlink(self.bc_path)
73        os.unlink(self.bc_path2)
74
75def test_main():
76    test_support.run_unittest(CompileallTests)
77
78
79if __name__ == "__main__":
80    test_main()
81