Issue #23491: Implement PEP 441: Improving Python Zip Application Support

Thanks to Paul Moore for the PEP and implementation.

1 files changed, 179 insertions, 0 deletions

diff --git a/Lib/zipapp.py b/Lib/zipapp.py
new file mode 100644
index 0000000..3b8f9bf
--- /dev/null
+++ b/Lib/zipapp.py
@@ -0,0 +1,179 @@
+import contextlib
+import os
+import pathlib
+import shutil
+import stat
+import sys
+import zipfile
+
+__all__ = ['ZipAppError', 'create_archive', 'get_interpreter']
+
+
+# The __main__.py used if the users specifies "-m module:fn".
+# Note that this will always be written as UTF-8 (module and
+# function names can be non-ASCII in Python 3).
+# We add a coding cookie even though UTF-8 is the default in Python 3
+# because the resulting archive may be intended to be run under Python 2.
+MAIN_TEMPLATE = """\
+# -*- coding: utf-8 -*-
+import {module}
+{module}.{fn}()
+"""
+
+
+# The Windows launcher defaults to UTF-8 when parsing shebang lines if the
+# file has no BOM. So use UTF-8 on Windows.
+# On Unix, use the filesystem encoding.
+if sys.platform.startswith('win'):
+    shebang_encoding = 'utf-8'
+else:
+    shebang_encoding = sys.getfilesystemencoding()
+
+
+class ZipAppError(ValueError):
+    pass
+
+
+@contextlib.contextmanager
+def _maybe_open(archive, mode):
+    if isinstance(archive, str):
+        with open(archive, mode) as f:
+            yield f
+    else:
+        yield archive
+
+
+def _write_file_prefix(f, interpreter):
+    """Write a shebang line."""
+    if interpreter:
+        shebang = b'#!%b\n' % (interpreter.encode(shebang_encoding),)
+        f.write(shebang)
+
+
+def _copy_archive(archive, new_archive, interpreter=None):
+    """Copy an application archive, modifying the shebang line."""
+    with _maybe_open(archive, 'rb') as src:
+        # Skip the shebang line from the source.
+        # Read 2 bytes of the source and check if they are #!.
+        first_2 = src.read(2)
+        if first_2 == b'#!':
+            # Discard the initial 2 bytes and the rest of the shebang line.
+            first_2 = b''
+            src.readline()
+
+        with _maybe_open(new_archive, 'wb') as dst:
+            _write_file_prefix(dst, interpreter)
+            # If there was no shebang, "first_2" contains the first 2 bytes
+            # of the source file, so write them before copying the rest
+            # of the file.
+            dst.write(first_2)
+            shutil.copyfileobj(src, dst)
+
+    if interpreter and isinstance(new_archive, str):
+        os.chmod(new_archive, os.stat(new_archive).st_mode | stat.S_IEXEC)
+
+
+def create_archive(source, target=None, interpreter=None, main=None):
+    """Create an application archive from SOURCE.
+
+    The SOURCE can be the name of a directory, or a filename or a file-like
+    object referring to an existing archive.
+
+    The content of SOURCE is packed into an application archive in TARGET,
+    which can be a filename or a file-like object.  If SOURCE is a directory,
+    TARGET can be omitted and will default to the name of SOURCE with .pyz
+    appended.
+
+    The created application archive will have a shebang line specifying
+    that it should run with INTERPRETER (there will be no shebang line if
+    INTERPRETER is None), and a __main__.py which runs MAIN (if MAIN is
+    not specified, an existing __main__.py will be used). It is an to specify
+    MAIN for anything other than a directory source with no __main__.py, and it
+    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)):
+        _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 main and has_main:
+        raise ZipAppError(
+            "Cannot specify entry point if the source has __main__.py")
+    if not (main or has_main):
+        raise ZipAppError("Archive has no entry point")
+
+    main_py = None
+    if main:
+        # Check that main has the right format
+        mod, sep, fn = main.partition(':')
+        mod_ok = all(part.isidentifier() for part in mod.split('.'))
+        fn_ok = all(part.isidentifier() for part in fn.split('.'))
+        if not (sep == ':' and mod_ok and fn_ok):
+            raise ZipAppError("Invalid entry point: " + main)
+        main_py = MAIN_TEMPLATE.format(module=mod, fn=fn)
+
+    if target is None:
+        target = source + '.pyz'
+
+    with _maybe_open(target, 'wb') as fd:
+        _write_file_prefix(fd, interpreter)
+        with zipfile.ZipFile(fd, 'w') as z:
+            root = pathlib.Path(source)
+            for child in root.rglob('*'):
+                arcname = str(child.relative_to(root))
+                z.write(str(child), arcname)
+            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)
+
+
+def get_interpreter(archive):
+    with _maybe_open(archive, 'rb') as f:
+        if f.read(2) == b'#!':
+            return f.readline().strip().decode(shebang_encoding)
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output', '-o', default=None,
+            help="The name of the output archive. "
+                 "Required if SOURCE is an archive.")
+    parser.add_argument('--python', '-p', default=None,
+            help="The name of the Python interpreter to use "
+                 "(default: no shebang line).")
+    parser.add_argument('--main', '-m', default=None,
+            help="The main function of the application "
+                 "(default: use an existing __main__.py).")
+    parser.add_argument('--info', default=False, action='store_true',
+            help="Display the interpreter from the archive.")
+    parser.add_argument('source',
+            help="Source directory (or existing archive).")
+
+    args = parser.parse_args()
+
+    # Handle `python -m zipapp archive.pyz --info`.
+    if args.info:
+        if not os.path.isfile(args.source):
+            raise SystemExit("Can only get info for an archive file")
+        interpreter = get_interpreter(args.source)
+        print("Interpreter: {}".format(interpreter or "<none>"))
+        sys.exit(0)
+
+    if os.path.isfile(args.source):
+        if args.output is None or 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")
+
+    create_archive(args.source, args.output,
+                   interpreter=args.python, main=args.main)
+
+
+if __name__ == '__main__':
+    main()