Python code coverage for Lib/zipapp.py

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