10a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport compileall
20a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport imp
30a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport os
40a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport py_compile
50a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport shutil
60a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport struct
70a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport tempfile
80a8c90248264a8b26970b4473770bcc3df8515fJosh Gaofrom test import test_support
90a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoimport unittest
100a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
110a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
120a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoclass CompileallTests(unittest.TestCase):
130a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
140a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def setUp(self):
150a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.directory = tempfile.mkdtemp()
160a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.source_path = os.path.join(self.directory, '_test.py')
170a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.bc_path = self.source_path + ('c' if __debug__ else 'o')
180a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        with open(self.source_path, 'w') as file:
190a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            file.write('x = 123\n')
200a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.source_path2 = os.path.join(self.directory, '_test2.py')
210a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.bc_path2 = self.source_path2 + ('c' if __debug__ else 'o')
220a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        shutil.copyfile(self.source_path, self.source_path2)
230a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
240a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def tearDown(self):
250a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        shutil.rmtree(self.directory)
260a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
270a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def data(self):
280a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        with open(self.bc_path, 'rb') as file:
290a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            data = file.read(8)
300a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        mtime = int(os.stat(self.source_path).st_mtime)
310a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        compare = struct.pack('<4sl', imp.get_magic(), mtime)
320a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        return data, compare
330a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
340a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def recreation_check(self, metadata):
350a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        """Check that compileall recreates bytecode when the new metadata is
360a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        used."""
370a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        if not hasattr(os, 'stat'):
380a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            return
390a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        py_compile.compile(self.source_path)
400a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertEqual(*self.data())
410a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        with open(self.bc_path, 'rb') as file:
420a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            bc = file.read()[len(metadata):]
430a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        with open(self.bc_path, 'wb') as file:
440a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            file.write(metadata)
450a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            file.write(bc)
460a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertNotEqual(*self.data())
470a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        compileall.compile_dir(self.directory, force=False, quiet=True)
480a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertTrue(*self.data())
490a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
500a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def test_mtime(self):
510a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # Test a change in mtime leads to a new .pyc.
520a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.recreation_check(struct.pack('<4sl', imp.get_magic(), 1))
530a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
540a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def test_magic_number(self):
550a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # Test a change in mtime leads to a new .pyc.
560a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.recreation_check(b'\0\0\0\0')
570a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
580a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    def test_compile_files(self):
590a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        # Test compiling a single file, and complete directory
600a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        for fn in (self.bc_path, self.bc_path2):
610a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            try:
620a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                os.unlink(fn)
630a8c90248264a8b26970b4473770bcc3df8515fJosh Gao            except:
640a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                pass
650a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        compileall.compile_file(self.source_path, force=False, quiet=True)
660a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertTrue(os.path.isfile(self.bc_path) \
670a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        and not os.path.isfile(self.bc_path2))
680a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        os.unlink(self.bc_path)
690a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        compileall.compile_dir(self.directory, force=False, quiet=True)
700a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        self.assertTrue(os.path.isfile(self.bc_path) \
710a8c90248264a8b26970b4473770bcc3df8515fJosh Gao                        and os.path.isfile(self.bc_path2))
720a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        os.unlink(self.bc_path)
730a8c90248264a8b26970b4473770bcc3df8515fJosh Gao        os.unlink(self.bc_path2)
740a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
750a8c90248264a8b26970b4473770bcc3df8515fJosh Gaodef test_main():
760a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    test_support.run_unittest(CompileallTests)
770a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
780a8c90248264a8b26970b4473770bcc3df8515fJosh Gao
790a8c90248264a8b26970b4473770bcc3df8515fJosh Gaoif __name__ == "__main__":
800a8c90248264a8b26970b4473770bcc3df8515fJosh Gao    test_main()
81