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, pathlib.Path):

archive = str(archive)

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'#!' + interpreter.encode(shebang_encoding) + b'\n'

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 error

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?

source_is_file = False

if hasattr(source, 'read') and hasattr(source, 'readline'):

source_is_file = True

else:

source = pathlib.Path(source)

if source.is_file():

source_is_file = True

if source_is_file:

_copy_archive(source, target, interpreter)

return

# We are creating a new archive from a directory.

if not source.exists():

raise ZipAppError("Source does not exist")

has_main = (source / '__main__.py').is_file()

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.with_suffix('.pyz')

elif not hasattr(target, 'write'):

target = pathlib.Path(target)

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 not hasattr(target, 'write'):

target.chmod(target.stat().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(args=None):

"""Run the zipapp command line interface.

The ARGS parameter lets you specify the argument list directly.

Omitting ARGS (or setting it to None) works as for argparse, using

sys.argv[1:] as the argument list.

"""

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

sys.exit(0)

if os.path.isfile(args.source):

if args.output is None or (os.path.exists(args.output) and

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