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 "")) 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()