• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import contextlib
2import os
3import pathlib
4import shutil
5import stat
6import sys
7import zipfile
8
9__all__ = ['ZipAppError', 'create_archive', 'get_interpreter']
10
11
12# The __main__.py used if the users specifies "-m module:fn".
13# Note that this will always be written as UTF-8 (module and
14# function names can be non-ASCII in Python 3).
15# We add a coding cookie even though UTF-8 is the default in Python 3
16# because the resulting archive may be intended to be run under Python 2.
17MAIN_TEMPLATE = """\
18# -*- coding: utf-8 -*-
19import {module}
20{module}.{fn}()
21"""
22
23
24# The Windows launcher defaults to UTF-8 when parsing shebang lines if the
25# file has no BOM. So use UTF-8 on Windows.
26# On Unix, use the filesystem encoding.
27if sys.platform.startswith('win'):
28    shebang_encoding = 'utf-8'
29else:
30    shebang_encoding = sys.getfilesystemencoding()
31
32
33class ZipAppError(ValueError):
34    pass
35
36
37@contextlib.contextmanager
38def _maybe_open(archive, mode):
39    if isinstance(archive, pathlib.Path):
40        archive = str(archive)
41    if isinstance(archive, str):
42        with open(archive, mode) as f:
43            yield f
44    else:
45        yield archive
46
47
48def _write_file_prefix(f, interpreter):
49    """Write a shebang line."""
50    if interpreter:
51        shebang = b'#!' + interpreter.encode(shebang_encoding) + b'\n'
52        f.write(shebang)
53
54
55def _copy_archive(archive, new_archive, interpreter=None):
56    """Copy an application archive, modifying the shebang line."""
57    with _maybe_open(archive, 'rb') as src:
58        # Skip the shebang line from the source.
59        # Read 2 bytes of the source and check if they are #!.
60        first_2 = src.read(2)
61        if first_2 == b'#!':
62            # Discard the initial 2 bytes and the rest of the shebang line.
63            first_2 = b''
64            src.readline()
65
66        with _maybe_open(new_archive, 'wb') as dst:
67            _write_file_prefix(dst, interpreter)
68            # If there was no shebang, "first_2" contains the first 2 bytes
69            # of the source file, so write them before copying the rest
70            # of the file.
71            dst.write(first_2)
72            shutil.copyfileobj(src, dst)
73
74    if interpreter and isinstance(new_archive, str):
75        os.chmod(new_archive, os.stat(new_archive).st_mode | stat.S_IEXEC)
76
77
78def create_archive(source, target=None, interpreter=None, main=None):
79    """Create an application archive from SOURCE.
80
81    The SOURCE can be the name of a directory, or a filename or a file-like
82    object referring to an existing archive.
83
84    The content of SOURCE is packed into an application archive in TARGET,
85    which can be a filename or a file-like object.  If SOURCE is a directory,
86    TARGET can be omitted and will default to the name of SOURCE with .pyz
87    appended.
88
89    The created application archive will have a shebang line specifying
90    that it should run with INTERPRETER (there will be no shebang line if
91    INTERPRETER is None), and a __main__.py which runs MAIN (if MAIN is
92    not specified, an existing __main__.py will be used).  It is an error
93    to specify MAIN for anything other than a directory source with no
94    __main__.py, and it is an error to omit MAIN if the directory has no
95    __main__.py.
96    """
97    # Are we copying an existing archive?
98    source_is_file = False
99    if hasattr(source, 'read') and hasattr(source, 'readline'):
100        source_is_file = True
101    else:
102        source = pathlib.Path(source)
103        if source.is_file():
104            source_is_file = True
105
106    if source_is_file:
107        _copy_archive(source, target, interpreter)
108        return
109
110    # We are creating a new archive from a directory.
111    if not source.exists():
112        raise ZipAppError("Source does not exist")
113    has_main = (source / '__main__.py').is_file()
114    if main and has_main:
115        raise ZipAppError(
116            "Cannot specify entry point if the source has __main__.py")
117    if not (main or has_main):
118        raise ZipAppError("Archive has no entry point")
119
120    main_py = None
121    if main:
122        # Check that main has the right format.
123        mod, sep, fn = main.partition(':')
124        mod_ok = all(part.isidentifier() for part in mod.split('.'))
125        fn_ok = all(part.isidentifier() for part in fn.split('.'))
126        if not (sep == ':' and mod_ok and fn_ok):
127            raise ZipAppError("Invalid entry point: " + main)
128        main_py = MAIN_TEMPLATE.format(module=mod, fn=fn)
129
130    if target is None:
131        target = source.with_suffix('.pyz')
132    elif not hasattr(target, 'write'):
133        target = pathlib.Path(target)
134
135    with _maybe_open(target, 'wb') as fd:
136        _write_file_prefix(fd, interpreter)
137        with zipfile.ZipFile(fd, 'w') as z:
138            root = pathlib.Path(source)
139            for child in root.rglob('*'):
140                arcname = str(child.relative_to(root))
141                z.write(str(child), arcname)
142            if main_py:
143                z.writestr('__main__.py', main_py.encode('utf-8'))
144
145    if interpreter and not hasattr(target, 'write'):
146        target.chmod(target.stat().st_mode | stat.S_IEXEC)
147
148
149def get_interpreter(archive):
150    with _maybe_open(archive, 'rb') as f:
151        if f.read(2) == b'#!':
152            return f.readline().strip().decode(shebang_encoding)
153
154
155def main(args=None):
156    """Run the zipapp command line interface.
157
158    The ARGS parameter lets you specify the argument list directly.
159    Omitting ARGS (or setting it to None) works as for argparse, using
160    sys.argv[1:] as the argument list.
161    """
162    import argparse
163
164    parser = argparse.ArgumentParser()
165    parser.add_argument('--output', '-o', default=None,
166            help="The name of the output archive. "
167                 "Required if SOURCE is an archive.")
168    parser.add_argument('--python', '-p', default=None,
169            help="The name of the Python interpreter to use "
170                 "(default: no shebang line).")
171    parser.add_argument('--main', '-m', default=None,
172            help="The main function of the application "
173                 "(default: use an existing __main__.py).")
174    parser.add_argument('--info', default=False, action='store_true',
175            help="Display the interpreter from the archive.")
176    parser.add_argument('source',
177            help="Source directory (or existing archive).")
178
179    args = parser.parse_args(args)
180
181    # Handle `python -m zipapp archive.pyz --info`.
182    if args.info:
183        if not os.path.isfile(args.source):
184            raise SystemExit("Can only get info for an archive file")
185        interpreter = get_interpreter(args.source)
186        print("Interpreter: {}".format(interpreter or "<none>"))
187        sys.exit(0)
188
189    if os.path.isfile(args.source):
190        if args.output is None or (os.path.exists(args.output) and
191                                   os.path.samefile(args.source, args.output)):
192            raise SystemExit("In-place editing of archives is not supported")
193        if args.main:
194            raise SystemExit("Cannot change the main function when copying")
195
196    create_archive(args.source, args.output,
197                   interpreter=args.python, main=args.main)
198
199
200if __name__ == '__main__':
201    main()
202