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