#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.

3 files changed, 138 insertions, 17 deletions

diff --git a/Doc/library/zipapp.rst b/Doc/library/zipapp.rst
index 00250b0..b4484c8 100644
--- a/Doc/library/zipapp.rst
+++ b/Doc/library/zipapp.rst
@@ -104,12 +104,13 @@ The module defines two convenience functions:
    Create an application archive from *source*.  The source can be any
    of the following:
 
-   * The name of a directory, in which case a new application archive
-     will be created from the content of that directory.
-   * The name of an existing application archive file, in which case the file is
-     copied to the target (modifying it to reflect the value given for the
-     *interpreter* argument).  The file name should include the ``.pyz``
-     extension, if required.
+   * The name of a directory, or a :class:`pathlib.Path` object referring
+     to a directory, in which case a new application archive will be
+     created from the content of that directory.
+   * The name of an existing application archive file, or a :class:`pathlib.Path`
+     object referring to such a file, in which case the file is copied to
+     the target (modifying it to reflect the value given for the *interpreter*
+     argument).  The file name should include the ``.pyz`` extension, if required.
    * A file object open for reading in bytes mode.  The content of the
      file should be an application archive, and the file object is
      assumed to be positioned at the start of the archive.
@@ -117,8 +118,8 @@ The module defines two convenience functions:
    The *target* argument determines where the resulting archive will be
    written:
 
-   * If it is the name of a file, the archive will be written to that
-     file.
+   * If it is the name of a file, or a :class:`pathlb.Path` object,
+     the archive will be written to that file.
    * If it is an open file object, the archive will be written to that
      file object, which must be open for writing in bytes mode.
    * If the target is omitted (or None), the source must be a directory
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")