zipapp.py 7.4 KB

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