#23657 Don't explicitly do an isinstance check for str in zipapp

As a result, explicitly support pathlib.Path objects as arguments.
Also added tests for the CLI interface.

2 files changed, 129 insertions, 9 deletions

diff --git a/Lib/test/test_zipapp.py b/Lib/test/test_zipapp.py
index e85c93b..9734380 100644
--- a/Lib/test/test_zipapp.py
+++ b/Lib/test/test_zipapp.py
@@ -9,6 +9,7 @@ import unittest
 import zipapp
 import zipfile
 
+from unittest.mock import patch
 
 class ZipAppTest(unittest.TestCase):
 
@@ -28,6 +29,15 @@ class ZipAppTest(unittest.TestCase):
         zipapp.create_archive(str(source), str(target))
         self.assertTrue(target.is_file())
 
+    def test_create_archive_with_pathlib(self):
+        # Test packing a directory using Path objects for source and target.
+        source = self.tmpdir / 'source'
+        source.mkdir()
+        (source / '__main__.py').touch()
+        target = self.tmpdir / 'source.pyz'
+        zipapp.create_archive(source, target)
+        self.assertTrue(target.is_file())
+
     def test_create_archive_with_subdirs(self):
         # Test packing a directory includes entries for subdirectories.
         source = self.tmpdir / 'source'
@@ -184,6 +194,18 @@ class ZipAppTest(unittest.TestCase):
         zipapp.create_archive(str(target), new_target, interpreter='python2.7')
         self.assertTrue(new_target.getvalue().startswith(b'#!python2.7\n'))
 
+    def test_read_from_pathobj(self):
+        # Test that we can copy an archive using an pathlib.Path object
+        # for the source.
+        source = self.tmpdir / 'source'
+        source.mkdir()
+        (source / '__main__.py').touch()
+        target1 = self.tmpdir / 'target1.pyz'
+        target2 = self.tmpdir / 'target2.pyz'
+        zipapp.create_archive(source, target1, interpreter='python')
+        zipapp.create_archive(target1, target2, interpreter='python2.7')
+        self.assertEqual(zipapp.get_interpreter(target2), 'python2.7')
+
     def test_read_from_fileobj(self):
         # Test that we can copy an archive using an open file object.
         source = self.tmpdir / 'source'
@@ -246,5 +268,82 @@ class ZipAppTest(unittest.TestCase):
         self.assertFalse(target.stat().st_mode & stat.S_IEXEC)
 
 
+class ZipAppCmdlineTest(unittest.TestCase):
+
+    """Test zipapp module command line API."""
+
+    def setUp(self):
+        tmpdir = tempfile.TemporaryDirectory()
+        self.addCleanup(tmpdir.cleanup)
+        self.tmpdir = pathlib.Path(tmpdir.name)
+
+    def make_archive(self):
+        # Test that an archive with no shebang line is not made executable.
+        source = self.tmpdir / 'source'
+        source.mkdir()
+        (source / '__main__.py').touch()
+        target = self.tmpdir / 'source.pyz'
+        zipapp.create_archive(source, target)
+        return target
+
+    def test_cmdline_create(self):
+        # Test the basic command line API.
+        source = self.tmpdir / 'source'
+        source.mkdir()
+        (source / '__main__.py').touch()
+        args = [str(source)]
+        zipapp.main(args)
+        target = source.with_suffix('.pyz')
+        self.assertTrue(target.is_file())
+
+    def test_cmdline_copy(self):
+        # Test copying an archive.
+        original = self.make_archive()
+        target = self.tmpdir / 'target.pyz'
+        args = [str(original), '-o', str(target)]
+        zipapp.main(args)
+        self.assertTrue(target.is_file())
+
+    def test_cmdline_copy_inplace(self):
+        # Test copying an archive in place fails.
+        original = self.make_archive()
+        target = self.tmpdir / 'target.pyz'
+        args = [str(original), '-o', str(original)]
+        with self.assertRaises(SystemExit) as cm:
+            zipapp.main(args)
+        # Program should exit with a non-zero returm code.
+        self.assertTrue(cm.exception.code)
+
+    def test_cmdline_copy_change_main(self):
+        # Test copying an archive doesn't allow changing __main__.py.
+        original = self.make_archive()
+        target = self.tmpdir / 'target.pyz'
+        args = [str(original), '-o', str(target), '-m', 'foo:bar']
+        with self.assertRaises(SystemExit) as cm:
+            zipapp.main(args)
+        # Program should exit with a non-zero returm code.
+        self.assertTrue(cm.exception.code)
+
+    @patch('sys.stdout', new_callable=io.StringIO)
+    def test_info_command(self, mock_stdout):
+        # Test the output of the info command.
+        target = self.make_archive()
+        args = [str(target), '--info']
+        with self.assertRaises(SystemExit) as cm:
+            zipapp.main(args)
+        # Program should exit with a zero returm code.
+        self.assertEqual(cm.exception.code, 0)
+        self.assertEqual(mock_stdout.getvalue(), "Interpreter: <none>\n")
+
+    def test_info_error(self):
+        # Test the info command fails when the archive does not exist.
+        target = self.tmpdir / 'dummy.pyz'
+        args = [str(target), '--info']
+        with self.assertRaises(SystemExit) as cm:
+            zipapp.main(args)
+        # Program should exit with a non-zero returm code.
+        self.assertTrue(cm.exception.code)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/zipapp.py b/Lib/zipapp.py
index c01c4ca..c8380bf 100644
--- a/Lib/zipapp.py
+++ b/Lib/zipapp.py
@@ -36,6 +36,8 @@ class ZipAppError(ValueError):
 
 @contextlib.contextmanager
 def _maybe_open(archive, mode):
+    if isinstance(archive, pathlib.Path):
+        archive = str(archive)
     if isinstance(archive, str):
         with open(archive, mode) as f:
             yield f
@@ -46,7 +48,7 @@ def _maybe_open(archive, mode):
 def _write_file_prefix(f, interpreter):
     """Write a shebang line."""
     if interpreter:
-        shebang = b'#!%b\n' % (interpreter.encode(shebang_encoding),)
+        shebang = b'#!' + interpreter.encode(shebang_encoding) + b'\n'
         f.write(shebang)
 
 
@@ -92,12 +94,22 @@ def create_archive(source, target=None, interpreter=None, main=None):
     is an error to omit MAIN if the directory has no __main__.py.
     """
     # Are we copying an existing archive?
-    if not (isinstance(source, str) and os.path.isdir(source)):
+    source_is_file = False
+    if hasattr(source, 'read') and hasattr(source, 'readline'):
+        source_is_file = True
+    else:
+        source = pathlib.Path(source)
+        if source.is_file():
+            source_is_file = True
+
+    if source_is_file:
         _copy_archive(source, target, interpreter)
         return
 
     # We are creating a new archive from a directory.
-    has_main = os.path.exists(os.path.join(source, '__main__.py'))
+    if not source.exists():
+        raise ZipAppError("Source does not exist")
+    has_main = (source / '__main__.py').is_file()
     if main and has_main:
         raise ZipAppError(
             "Cannot specify entry point if the source has __main__.py")
@@ -115,7 +127,9 @@ def create_archive(source, target=None, interpreter=None, main=None):
         main_py = MAIN_TEMPLATE.format(module=mod, fn=fn)
 
     if target is None:
-        target = source + '.pyz'
+        target = source.with_suffix('.pyz')
+    elif not hasattr(target, 'write'):
+        target = pathlib.Path(target)
 
     with _maybe_open(target, 'wb') as fd:
         _write_file_prefix(fd, interpreter)
@@ -127,8 +141,8 @@ def create_archive(source, target=None, interpreter=None, main=None):
             if main_py:
                 z.writestr('__main__.py', main_py.encode('utf-8'))
 
-    if interpreter and isinstance(target, str):
-        os.chmod(target, os.stat(target).st_mode | stat.S_IEXEC)
+    if interpreter and not hasattr(target, 'write'):
+        target.chmod(target.stat().st_mode | stat.S_IEXEC)
 
 
 def get_interpreter(archive):
@@ -137,7 +151,13 @@ def get_interpreter(archive):
             return f.readline().strip().decode(shebang_encoding)
 
 
-def main():
+def main(args=None):
+    """Run the zipapp command line interface.
+
+    The ARGS parameter lets you specify the argument list directly.
+    Omitting ARGS (or setting it to None) works as for argparse, using
+    sys.argv[1:] as the argument list.
+    """
     import argparse
 
     parser = argparse.ArgumentParser()
@@ -155,7 +175,7 @@ def main():
     parser.add_argument('source',
             help="Source directory (or existing archive).")
 
-    args = parser.parse_args()
+    args = parser.parse_args(args)
 
     # Handle `python -m zipapp archive.pyz --info`.
     if args.info:
@@ -166,7 +186,8 @@ def main():
         sys.exit(0)
 
     if os.path.isfile(args.source):
-        if args.output is None or os.path.samefile(args.source, args.output):
+        if args.output is None or (os.path.exists(args.output) and
+                                   os.path.samefile(args.source, args.output)):
             raise SystemExit("In-place editing of archives is not supported")
         if args.main:
             raise SystemExit("Cannot change the main function when copying")