1"""Test harness for the zipapp module."""
2
3import io
4import pathlib
5import stat
6import sys
7import tempfile
8import unittest
9import zipapp
10import zipfile
11from test.support import requires_zlib
12from test.support import os_helper
13
14from unittest.mock import patch
15
16class ZipAppTest(unittest.TestCase):
17
18    """Test zipapp module functionality."""
19
20    def setUp(self):
21        tmpdir = tempfile.TemporaryDirectory()
22        self.addCleanup(tmpdir.cleanup)
23        self.tmpdir = pathlib.Path(tmpdir.name)
24
25    def test_create_archive(self):
26        # Test packing a directory.
27        source = self.tmpdir / 'source'
28        source.mkdir()
29        (source / '__main__.py').touch()
30        target = self.tmpdir / 'source.pyz'
31        zipapp.create_archive(str(source), str(target))
32        self.assertTrue(target.is_file())
33
34    def test_create_archive_with_pathlib(self):
35        # Test packing a directory using Path objects for source and target.
36        source = self.tmpdir / 'source'
37        source.mkdir()
38        (source / '__main__.py').touch()
39        target = self.tmpdir / 'source.pyz'
40        zipapp.create_archive(source, target)
41        self.assertTrue(target.is_file())
42
43    def test_create_archive_with_subdirs(self):
44        # Test packing a directory includes entries for subdirectories.
45        source = self.tmpdir / 'source'
46        source.mkdir()
47        (source / '__main__.py').touch()
48        (source / 'foo').mkdir()
49        (source / 'bar').mkdir()
50        (source / 'foo' / '__init__.py').touch()
51        target = io.BytesIO()
52        zipapp.create_archive(str(source), target)
53        target.seek(0)
54        with zipfile.ZipFile(target, 'r') as z:
55            self.assertIn('foo/', z.namelist())
56            self.assertIn('bar/', z.namelist())
57
58    def test_create_archive_with_filter(self):
59        # Test packing a directory and using filter to specify
60        # which files to include.
61        def skip_pyc_files(path):
62            return path.suffix != '.pyc'
63        source = self.tmpdir / 'source'
64        source.mkdir()
65        (source / '__main__.py').touch()
66        (source / 'test.py').touch()
67        (source / 'test.pyc').touch()
68        target = self.tmpdir / 'source.pyz'
69
70        zipapp.create_archive(source, target, filter=skip_pyc_files)
71        with zipfile.ZipFile(target, 'r') as z:
72            self.assertIn('__main__.py', z.namelist())
73            self.assertIn('test.py', z.namelist())
74            self.assertNotIn('test.pyc', z.namelist())
75
76    def test_create_archive_filter_exclude_dir(self):
77        # Test packing a directory and using a filter to exclude a
78        # subdirectory (ensures that the path supplied to include
79        # is relative to the source location, as expected).
80        def skip_dummy_dir(path):
81            return path.parts[0] != 'dummy'
82        source = self.tmpdir / 'source'
83        source.mkdir()
84        (source / '__main__.py').touch()
85        (source / 'test.py').touch()
86        (source / 'dummy').mkdir()
87        (source / 'dummy' / 'test2.py').touch()
88        target = self.tmpdir / 'source.pyz'
89
90        zipapp.create_archive(source, target, filter=skip_dummy_dir)
91        with zipfile.ZipFile(target, 'r') as z:
92            self.assertEqual(len(z.namelist()), 2)
93            self.assertIn('__main__.py', z.namelist())
94            self.assertIn('test.py', z.namelist())
95
96    def test_create_archive_default_target(self):
97        # Test packing a directory to the default name.
98        source = self.tmpdir / 'source'
99        source.mkdir()
100        (source / '__main__.py').touch()
101        zipapp.create_archive(str(source))
102        expected_target = self.tmpdir / 'source.pyz'
103        self.assertTrue(expected_target.is_file())
104
105    @requires_zlib()
106    def test_create_archive_with_compression(self):
107        # Test packing a directory into a compressed archive.
108        source = self.tmpdir / 'source'
109        source.mkdir()
110        (source / '__main__.py').touch()
111        (source / 'test.py').touch()
112        target = self.tmpdir / 'source.pyz'
113
114        zipapp.create_archive(source, target, compressed=True)
115        with zipfile.ZipFile(target, 'r') as z:
116            for name in ('__main__.py', 'test.py'):
117                self.assertEqual(z.getinfo(name).compress_type,
118                                 zipfile.ZIP_DEFLATED)
119
120    def test_no_main(self):
121        # Test that packing a directory with no __main__.py fails.
122        source = self.tmpdir / 'source'
123        source.mkdir()
124        (source / 'foo.py').touch()
125        target = self.tmpdir / 'source.pyz'
126        with self.assertRaises(zipapp.ZipAppError):
127            zipapp.create_archive(str(source), str(target))
128
129    def test_main_and_main_py(self):
130        # Test that supplying a main argument with __main__.py fails.
131        source = self.tmpdir / 'source'
132        source.mkdir()
133        (source / '__main__.py').touch()
134        target = self.tmpdir / 'source.pyz'
135        with self.assertRaises(zipapp.ZipAppError):
136            zipapp.create_archive(str(source), str(target), main='pkg.mod:fn')
137
138    def test_main_written(self):
139        # Test that the __main__.py is written correctly.
140        source = self.tmpdir / 'source'
141        source.mkdir()
142        (source / 'foo.py').touch()
143        target = self.tmpdir / 'source.pyz'
144        zipapp.create_archive(str(source), str(target), main='pkg.mod:fn')
145        with zipfile.ZipFile(str(target), 'r') as z:
146            self.assertIn('__main__.py', z.namelist())
147            self.assertIn(b'pkg.mod.fn()', z.read('__main__.py'))
148
149    def test_main_only_written_once(self):
150        # Test that we don't write multiple __main__.py files.
151        # The initial implementation had this bug; zip files allow
152        # multiple entries with the same name
153        source = self.tmpdir / 'source'
154        source.mkdir()
155        # Write 2 files, as the original bug wrote __main__.py
156        # once for each file written :-(
157        # See http://bugs.python.org/review/23491/diff/13982/Lib/zipapp.py#newcode67Lib/zipapp.py:67
158        # (line 67)
159        (source / 'foo.py').touch()
160        (source / 'bar.py').touch()
161        target = self.tmpdir / 'source.pyz'
162        zipapp.create_archive(str(source), str(target), main='pkg.mod:fn')
163        with zipfile.ZipFile(str(target), 'r') as z:
164            self.assertEqual(1, z.namelist().count('__main__.py'))
165
166    def test_main_validation(self):
167        # Test that invalid values for main are rejected.
168        source = self.tmpdir / 'source'
169        source.mkdir()
170        target = self.tmpdir / 'source.pyz'
171        problems = [
172            '', 'foo', 'foo:', ':bar', '12:bar', 'a.b.c.:d',
173            '.a:b', 'a:b.', 'a:.b', 'a:silly name'
174        ]
175        for main in problems:
176            with self.subTest(main=main):
177                with self.assertRaises(zipapp.ZipAppError):
178                    zipapp.create_archive(str(source), str(target), main=main)
179
180    def test_default_no_shebang(self):
181        # Test that no shebang line is written to the target by default.
182        source = self.tmpdir / 'source'
183        source.mkdir()
184        (source / '__main__.py').touch()
185        target = self.tmpdir / 'source.pyz'
186        zipapp.create_archive(str(source), str(target))
187        with target.open('rb') as f:
188            self.assertNotEqual(f.read(2), b'#!')
189
190    def test_custom_interpreter(self):
191        # Test that a shebang line with a custom interpreter is written
192        # correctly.
193        source = self.tmpdir / 'source'
194        source.mkdir()
195        (source / '__main__.py').touch()
196        target = self.tmpdir / 'source.pyz'
197        zipapp.create_archive(str(source), str(target), interpreter='python')
198        with target.open('rb') as f:
199            self.assertEqual(f.read(2), b'#!')
200            self.assertEqual(b'python\n', f.readline())
201
202    def test_pack_to_fileobj(self):
203        # Test that we can pack to a file object.
204        source = self.tmpdir / 'source'
205        source.mkdir()
206        (source / '__main__.py').touch()
207        target = io.BytesIO()
208        zipapp.create_archive(str(source), target, interpreter='python')
209        self.assertTrue(target.getvalue().startswith(b'#!python\n'))
210
211    def test_read_shebang(self):
212        # Test that we can read the shebang line correctly.
213        source = self.tmpdir / 'source'
214        source.mkdir()
215        (source / '__main__.py').touch()
216        target = self.tmpdir / 'source.pyz'
217        zipapp.create_archive(str(source), str(target), interpreter='python')
218        self.assertEqual(zipapp.get_interpreter(str(target)), 'python')
219
220    def test_read_missing_shebang(self):
221        # Test that reading the shebang line of a file without one returns None.
222        source = self.tmpdir / 'source'
223        source.mkdir()
224        (source / '__main__.py').touch()
225        target = self.tmpdir / 'source.pyz'
226        zipapp.create_archive(str(source), str(target))
227        self.assertEqual(zipapp.get_interpreter(str(target)), None)
228
229    def test_modify_shebang(self):
230        # Test that we can change the shebang of a file.
231        source = self.tmpdir / 'source'
232        source.mkdir()
233        (source / '__main__.py').touch()
234        target = self.tmpdir / 'source.pyz'
235        zipapp.create_archive(str(source), str(target), interpreter='python')
236        new_target = self.tmpdir / 'changed.pyz'
237        zipapp.create_archive(str(target), str(new_target), interpreter='python2.7')
238        self.assertEqual(zipapp.get_interpreter(str(new_target)), 'python2.7')
239
240    def test_write_shebang_to_fileobj(self):
241        # Test that we can change the shebang of a file, writing the result to a
242        # file object.
243        source = self.tmpdir / 'source'
244        source.mkdir()
245        (source / '__main__.py').touch()
246        target = self.tmpdir / 'source.pyz'
247        zipapp.create_archive(str(source), str(target), interpreter='python')
248        new_target = io.BytesIO()
249        zipapp.create_archive(str(target), new_target, interpreter='python2.7')
250        self.assertTrue(new_target.getvalue().startswith(b'#!python2.7\n'))
251
252    def test_read_from_pathobj(self):
253        # Test that we can copy an archive using a pathlib.Path object
254        # for the source.
255        source = self.tmpdir / 'source'
256        source.mkdir()
257        (source / '__main__.py').touch()
258        target1 = self.tmpdir / 'target1.pyz'
259        target2 = self.tmpdir / 'target2.pyz'
260        zipapp.create_archive(source, target1, interpreter='python')
261        zipapp.create_archive(target1, target2, interpreter='python2.7')
262        self.assertEqual(zipapp.get_interpreter(target2), 'python2.7')
263
264    def test_read_from_fileobj(self):
265        # Test that we can copy an archive using an open file object.
266        source = self.tmpdir / 'source'
267        source.mkdir()
268        (source / '__main__.py').touch()
269        target = self.tmpdir / 'source.pyz'
270        temp_archive = io.BytesIO()
271        zipapp.create_archive(str(source), temp_archive, interpreter='python')
272        new_target = io.BytesIO()
273        temp_archive.seek(0)
274        zipapp.create_archive(temp_archive, new_target, interpreter='python2.7')
275        self.assertTrue(new_target.getvalue().startswith(b'#!python2.7\n'))
276
277    def test_remove_shebang(self):
278        # Test that we can remove the shebang from a file.
279        source = self.tmpdir / 'source'
280        source.mkdir()
281        (source / '__main__.py').touch()
282        target = self.tmpdir / 'source.pyz'
283        zipapp.create_archive(str(source), str(target), interpreter='python')
284        new_target = self.tmpdir / 'changed.pyz'
285        zipapp.create_archive(str(target), str(new_target), interpreter=None)
286        self.assertEqual(zipapp.get_interpreter(str(new_target)), None)
287
288    def test_content_of_copied_archive(self):
289        # Test that copying an archive doesn't corrupt it.
290        source = self.tmpdir / 'source'
291        source.mkdir()
292        (source / '__main__.py').touch()
293        target = io.BytesIO()
294        zipapp.create_archive(str(source), target, interpreter='python')
295        new_target = io.BytesIO()
296        target.seek(0)
297        zipapp.create_archive(target, new_target, interpreter=None)
298        new_target.seek(0)
299        with zipfile.ZipFile(new_target, 'r') as z:
300            self.assertEqual(set(z.namelist()), {'__main__.py'})
301
302    # (Unix only) tests that archives with shebang lines are made executable
303    @unittest.skipIf(sys.platform == 'win32',
304                     'Windows does not support an executable bit')
305    @os_helper.skip_unless_working_chmod
306    def test_shebang_is_executable(self):
307        # Test that an archive with a shebang line is made executable.
308        source = self.tmpdir / 'source'
309        source.mkdir()
310        (source / '__main__.py').touch()
311        target = self.tmpdir / 'source.pyz'
312        zipapp.create_archive(str(source), str(target), interpreter='python')
313        self.assertTrue(target.stat().st_mode & stat.S_IEXEC)
314
315    @unittest.skipIf(sys.platform == 'win32',
316                     'Windows does not support an executable bit')
317    def test_no_shebang_is_not_executable(self):
318        # Test that an archive with no shebang line is not made executable.
319        source = self.tmpdir / 'source'
320        source.mkdir()
321        (source / '__main__.py').touch()
322        target = self.tmpdir / 'source.pyz'
323        zipapp.create_archive(str(source), str(target), interpreter=None)
324        self.assertFalse(target.stat().st_mode & stat.S_IEXEC)
325
326
327class ZipAppCmdlineTest(unittest.TestCase):
328
329    """Test zipapp module command line API."""
330
331    def setUp(self):
332        tmpdir = tempfile.TemporaryDirectory()
333        self.addCleanup(tmpdir.cleanup)
334        self.tmpdir = pathlib.Path(tmpdir.name)
335
336    def make_archive(self):
337        # Test that an archive with no shebang line is not made executable.
338        source = self.tmpdir / 'source'
339        source.mkdir()
340        (source / '__main__.py').touch()
341        target = self.tmpdir / 'source.pyz'
342        zipapp.create_archive(source, target)
343        return target
344
345    def test_cmdline_create(self):
346        # Test the basic command line API.
347        source = self.tmpdir / 'source'
348        source.mkdir()
349        (source / '__main__.py').touch()
350        args = [str(source)]
351        zipapp.main(args)
352        target = source.with_suffix('.pyz')
353        self.assertTrue(target.is_file())
354
355    def test_cmdline_copy(self):
356        # Test copying an archive.
357        original = self.make_archive()
358        target = self.tmpdir / 'target.pyz'
359        args = [str(original), '-o', str(target)]
360        zipapp.main(args)
361        self.assertTrue(target.is_file())
362
363    def test_cmdline_copy_inplace(self):
364        # Test copying an archive in place fails.
365        original = self.make_archive()
366        target = self.tmpdir / 'target.pyz'
367        args = [str(original), '-o', str(original)]
368        with self.assertRaises(SystemExit) as cm:
369            zipapp.main(args)
370        # Program should exit with a non-zero return code.
371        self.assertTrue(cm.exception.code)
372
373    def test_cmdline_copy_change_main(self):
374        # Test copying an archive doesn't allow changing __main__.py.
375        original = self.make_archive()
376        target = self.tmpdir / 'target.pyz'
377        args = [str(original), '-o', str(target), '-m', 'foo:bar']
378        with self.assertRaises(SystemExit) as cm:
379            zipapp.main(args)
380        # Program should exit with a non-zero return code.
381        self.assertTrue(cm.exception.code)
382
383    @patch('sys.stdout', new_callable=io.StringIO)
384    def test_info_command(self, mock_stdout):
385        # Test the output of the info command.
386        target = self.make_archive()
387        args = [str(target), '--info']
388        with self.assertRaises(SystemExit) as cm:
389            zipapp.main(args)
390        # Program should exit with a zero return code.
391        self.assertEqual(cm.exception.code, 0)
392        self.assertEqual(mock_stdout.getvalue(), "Interpreter: <none>\n")
393
394    def test_info_error(self):
395        # Test the info command fails when the archive does not exist.
396        target = self.tmpdir / 'dummy.pyz'
397        args = [str(target), '--info']
398        with self.assertRaises(SystemExit) as cm:
399            zipapp.main(args)
400        # Program should exit with a non-zero return code.
401        self.assertTrue(cm.exception.code)
402
403
404if __name__ == "__main__":
405    unittest.main()
406