Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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 ZipAppError(ValueError): |
| 34 | pass |
| 35 | |
| 36 | |
| 37 | @contextlib.contextmanager |
| 38 | def _maybe_open(archive, mode): |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 39 | if isinstance(archive, pathlib.Path): |
| 40 | archive = str(archive) |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 41 | if isinstance(archive, str): |
| 42 | with open(archive, mode) as f: |
| 43 | yield f |
| 44 | else: |
| 45 | yield archive |
| 46 | |
| 47 | |
| 48 | def _write_file_prefix(f, interpreter): |
| 49 | """Write a shebang line.""" |
| 50 | if interpreter: |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 51 | shebang = b'#!' + interpreter.encode(shebang_encoding) + b'\n' |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 52 | f.write(shebang) |
| 53 | |
| 54 | |
| 55 | def _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 | |
| 78 | def 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 |
Serhiy Storchaka | 6a7b3a7 | 2016-04-17 08:32:47 +0300 | [diff] [blame] | 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. |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 96 | """ |
| 97 | # Are we copying an existing archive? |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 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: |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 107 | _copy_archive(source, target, interpreter) |
| 108 | return |
| 109 | |
Brett Cannon | 469d0fb | 2015-03-13 11:13:20 -0400 | [diff] [blame] | 110 | # We are creating a new archive from a directory. |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 111 | if not source.exists(): |
| 112 | raise ZipAppError("Source does not exist") |
| 113 | has_main = (source / '__main__.py').is_file() |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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: |
Brett Cannon | 469d0fb | 2015-03-13 11:13:20 -0400 | [diff] [blame] | 122 | # Check that main has the right format. |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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: |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 131 | target = source.with_suffix('.pyz') |
| 132 | elif not hasattr(target, 'write'): |
| 133 | target = pathlib.Path(target) |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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 | |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 145 | if interpreter and not hasattr(target, 'write'): |
| 146 | target.chmod(target.stat().st_mode | stat.S_IEXEC) |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 147 | |
| 148 | |
| 149 | def 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 | |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 155 | def 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 | """ |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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 | |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 179 | args = parser.parse_args(args) |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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): |
Paul Moore | a4d4dd3 | 2015-03-22 15:32:36 +0000 | [diff] [blame] | 190 | if args.output is None or (os.path.exists(args.output) and |
| 191 | os.path.samefile(args.source, args.output)): |
Brett Cannon | cc4dfc1 | 2015-03-13 10:40:49 -0400 | [diff] [blame] | 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 | |
| 200 | if __name__ == '__main__': |
| 201 | main() |