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