1import functools
2import importlib.util
3import os
4import py_compile
5import shutil
6import stat
7import subprocess
8import sys
9import tempfile
10import unittest
11
12from test import support
13from test.support import os_helper, script_helper
14
15
16def without_source_date_epoch(fxn):
17    """Runs function with SOURCE_DATE_EPOCH unset."""
18    @functools.wraps(fxn)
19    def wrapper(*args, **kwargs):
20        with os_helper.EnvironmentVarGuard() as env:
21            env.unset('SOURCE_DATE_EPOCH')
22            return fxn(*args, **kwargs)
23    return wrapper
24
25
26def with_source_date_epoch(fxn):
27    """Runs function with SOURCE_DATE_EPOCH set."""
28    @functools.wraps(fxn)
29    def wrapper(*args, **kwargs):
30        with os_helper.EnvironmentVarGuard() as env:
31            env['SOURCE_DATE_EPOCH'] = '123456789'
32            return fxn(*args, **kwargs)
33    return wrapper
34
35
36# Run tests with SOURCE_DATE_EPOCH set or unset explicitly.
37class SourceDateEpochTestMeta(type(unittest.TestCase)):
38    def __new__(mcls, name, bases, dct, *, source_date_epoch):
39        cls = super().__new__(mcls, name, bases, dct)
40
41        for attr in dir(cls):
42            if attr.startswith('test_'):
43                meth = getattr(cls, attr)
44                if source_date_epoch:
45                    wrapper = with_source_date_epoch(meth)
46                else:
47                    wrapper = without_source_date_epoch(meth)
48                setattr(cls, attr, wrapper)
49
50        return cls
51
52
53class PyCompileTestsBase:
54
55    def setUp(self):
56        self.directory = tempfile.mkdtemp(dir=os.getcwd())
57        self.source_path = os.path.join(self.directory, '_test.py')
58        self.pyc_path = self.source_path + 'c'
59        self.cache_path = importlib.util.cache_from_source(self.source_path)
60        self.cwd_drive = os.path.splitdrive(os.getcwd())[0]
61        # In these tests we compute relative paths.  When using Windows, the
62        # current working directory path and the 'self.source_path' might be
63        # on different drives.  Therefore we need to switch to the drive where
64        # the temporary source file lives.
65        drive = os.path.splitdrive(self.source_path)[0]
66        if drive:
67            os.chdir(drive)
68        with open(self.source_path, 'w') as file:
69            file.write('x = 123\n')
70
71    def tearDown(self):
72        shutil.rmtree(self.directory)
73        if self.cwd_drive:
74            os.chdir(self.cwd_drive)
75
76    def test_absolute_path(self):
77        py_compile.compile(self.source_path, self.pyc_path)
78        self.assertTrue(os.path.exists(self.pyc_path))
79        self.assertFalse(os.path.exists(self.cache_path))
80
81    def test_do_not_overwrite_symlinks(self):
82        # In the face of a cfile argument being a symlink, bail out.
83        # Issue #17222
84        try:
85            os.symlink(self.pyc_path + '.actual', self.pyc_path)
86        except (NotImplementedError, OSError):
87            self.skipTest('need to be able to create a symlink for a file')
88        else:
89            assert os.path.islink(self.pyc_path)
90            with self.assertRaises(FileExistsError):
91                py_compile.compile(self.source_path, self.pyc_path)
92
93    @unittest.skipIf(not os.path.exists(os.devnull) or os.path.isfile(os.devnull),
94                     'requires os.devnull and for it to be a non-regular file')
95    def test_do_not_overwrite_nonregular_files(self):
96        # In the face of a cfile argument being a non-regular file, bail out.
97        # Issue #17222
98        with self.assertRaises(FileExistsError):
99            py_compile.compile(self.source_path, os.devnull)
100
101    def test_cache_path(self):
102        py_compile.compile(self.source_path)
103        self.assertTrue(os.path.exists(self.cache_path))
104
105    def test_cwd(self):
106        with os_helper.change_cwd(self.directory):
107            py_compile.compile(os.path.basename(self.source_path),
108                               os.path.basename(self.pyc_path))
109        self.assertTrue(os.path.exists(self.pyc_path))
110        self.assertFalse(os.path.exists(self.cache_path))
111
112    def test_relative_path(self):
113        py_compile.compile(os.path.relpath(self.source_path),
114                           os.path.relpath(self.pyc_path))
115        self.assertTrue(os.path.exists(self.pyc_path))
116        self.assertFalse(os.path.exists(self.cache_path))
117
118    @unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,
119                     'non-root user required')
120    @unittest.skipIf(os.name == 'nt',
121                     'cannot control directory permissions on Windows')
122    def test_exceptions_propagate(self):
123        # Make sure that exceptions raised thanks to issues with writing
124        # bytecode.
125        # http://bugs.python.org/issue17244
126        mode = os.stat(self.directory)
127        os.chmod(self.directory, stat.S_IREAD)
128        try:
129            with self.assertRaises(IOError):
130                py_compile.compile(self.source_path, self.pyc_path)
131        finally:
132            os.chmod(self.directory, mode.st_mode)
133
134    def test_bad_coding(self):
135        bad_coding = os.path.join(os.path.dirname(__file__), 'bad_coding2.py')
136        with support.captured_stderr():
137            self.assertIsNone(py_compile.compile(bad_coding, doraise=False))
138        self.assertFalse(os.path.exists(
139            importlib.util.cache_from_source(bad_coding)))
140
141    def test_source_date_epoch(self):
142        py_compile.compile(self.source_path, self.pyc_path)
143        self.assertTrue(os.path.exists(self.pyc_path))
144        self.assertFalse(os.path.exists(self.cache_path))
145        with open(self.pyc_path, 'rb') as fp:
146            flags = importlib._bootstrap_external._classify_pyc(
147                fp.read(), 'test', {})
148        if os.environ.get('SOURCE_DATE_EPOCH'):
149            expected_flags = 0b11
150        else:
151            expected_flags = 0b00
152
153        self.assertEqual(flags, expected_flags)
154
155    @unittest.skipIf(sys.flags.optimize > 0, 'test does not work with -O')
156    def test_double_dot_no_clobber(self):
157        # http://bugs.python.org/issue22966
158        # py_compile foo.bar.py -> __pycache__/foo.cpython-34.pyc
159        weird_path = os.path.join(self.directory, 'foo.bar.py')
160        cache_path = importlib.util.cache_from_source(weird_path)
161        pyc_path = weird_path + 'c'
162        head, tail = os.path.split(cache_path)
163        penultimate_tail = os.path.basename(head)
164        self.assertEqual(
165            os.path.join(penultimate_tail, tail),
166            os.path.join(
167                '__pycache__',
168                'foo.bar.{}.pyc'.format(sys.implementation.cache_tag)))
169        with open(weird_path, 'w') as file:
170            file.write('x = 123\n')
171        py_compile.compile(weird_path)
172        self.assertTrue(os.path.exists(cache_path))
173        self.assertFalse(os.path.exists(pyc_path))
174
175    def test_optimization_path(self):
176        # Specifying optimized bytecode should lead to a path reflecting that.
177        self.assertIn('opt-2', py_compile.compile(self.source_path, optimize=2))
178
179    def test_invalidation_mode(self):
180        py_compile.compile(
181            self.source_path,
182            invalidation_mode=py_compile.PycInvalidationMode.CHECKED_HASH,
183        )
184        with open(self.cache_path, 'rb') as fp:
185            flags = importlib._bootstrap_external._classify_pyc(
186                fp.read(), 'test', {})
187        self.assertEqual(flags, 0b11)
188        py_compile.compile(
189            self.source_path,
190            invalidation_mode=py_compile.PycInvalidationMode.UNCHECKED_HASH,
191        )
192        with open(self.cache_path, 'rb') as fp:
193            flags = importlib._bootstrap_external._classify_pyc(
194                fp.read(), 'test', {})
195        self.assertEqual(flags, 0b1)
196
197    def test_quiet(self):
198        bad_coding = os.path.join(os.path.dirname(__file__), 'bad_coding2.py')
199        with support.captured_stderr() as stderr:
200            self.assertIsNone(py_compile.compile(bad_coding, doraise=False, quiet=2))
201            self.assertIsNone(py_compile.compile(bad_coding, doraise=True, quiet=2))
202            self.assertEqual(stderr.getvalue(), '')
203            with self.assertRaises(py_compile.PyCompileError):
204                py_compile.compile(bad_coding, doraise=True, quiet=1)
205
206
207class PyCompileTestsWithSourceEpoch(PyCompileTestsBase,
208                                    unittest.TestCase,
209                                    metaclass=SourceDateEpochTestMeta,
210                                    source_date_epoch=True):
211    pass
212
213
214class PyCompileTestsWithoutSourceEpoch(PyCompileTestsBase,
215                                       unittest.TestCase,
216                                       metaclass=SourceDateEpochTestMeta,
217                                       source_date_epoch=False):
218    pass
219
220
221class PyCompileCLITestCase(unittest.TestCase):
222
223    def setUp(self):
224        self.directory = tempfile.mkdtemp()
225        self.source_path = os.path.join(self.directory, '_test.py')
226        self.cache_path = importlib.util.cache_from_source(self.source_path)
227        with open(self.source_path, 'w') as file:
228            file.write('x = 123\n')
229
230    def tearDown(self):
231        os_helper.rmtree(self.directory)
232
233    def pycompilecmd(self, *args, **kwargs):
234        # assert_python_* helpers don't return proc object. We'll just use
235        # subprocess.run() instead of spawn_python() and its friends to test
236        # stdin support of the CLI.
237        if args and args[0] == '-' and 'input' in kwargs:
238            return subprocess.run([sys.executable, '-m', 'py_compile', '-'],
239                                  input=kwargs['input'].encode(),
240                                  capture_output=True)
241        return script_helper.assert_python_ok('-m', 'py_compile', *args, **kwargs)
242
243    def pycompilecmd_failure(self, *args):
244        return script_helper.assert_python_failure('-m', 'py_compile', *args)
245
246    def test_stdin(self):
247        result = self.pycompilecmd('-', input=self.source_path)
248        self.assertEqual(result.returncode, 0)
249        self.assertEqual(result.stdout, b'')
250        self.assertEqual(result.stderr, b'')
251        self.assertTrue(os.path.exists(self.cache_path))
252
253    def test_with_files(self):
254        rc, stdout, stderr = self.pycompilecmd(self.source_path, self.source_path)
255        self.assertEqual(rc, 0)
256        self.assertEqual(stdout, b'')
257        self.assertEqual(stderr, b'')
258        self.assertTrue(os.path.exists(self.cache_path))
259
260    def test_bad_syntax(self):
261        bad_syntax = os.path.join(os.path.dirname(__file__), 'badsyntax_3131.py')
262        rc, stdout, stderr = self.pycompilecmd_failure(bad_syntax)
263        self.assertEqual(rc, 1)
264        self.assertEqual(stdout, b'')
265        self.assertIn(b'SyntaxError', stderr)
266
267    def test_bad_syntax_with_quiet(self):
268        bad_syntax = os.path.join(os.path.dirname(__file__), 'badsyntax_3131.py')
269        rc, stdout, stderr = self.pycompilecmd_failure('-q', bad_syntax)
270        self.assertEqual(rc, 1)
271        self.assertEqual(stdout, b'')
272        self.assertEqual(stderr, b'')
273
274    def test_file_not_exists(self):
275        should_not_exists = os.path.join(os.path.dirname(__file__), 'should_not_exists.py')
276        rc, stdout, stderr = self.pycompilecmd_failure(self.source_path, should_not_exists)
277        self.assertEqual(rc, 1)
278        self.assertEqual(stdout, b'')
279        self.assertIn(b'no such file or directory', stderr.lower())
280
281    def test_file_not_exists_with_quiet(self):
282        should_not_exists = os.path.join(os.path.dirname(__file__), 'should_not_exists.py')
283        rc, stdout, stderr = self.pycompilecmd_failure('-q', self.source_path, should_not_exists)
284        self.assertEqual(rc, 1)
285        self.assertEqual(stdout, b'')
286        self.assertEqual(stderr, b'')
287
288
289if __name__ == "__main__":
290    unittest.main()
291