taisei/scripts/pack.py
2021-03-28 14:43:02 +03:00

189 lines
4.9 KiB
Python
Executable file

#!/usr/bin/env python3
import os
import sys
import re
import zstandard
import zlib
import shutil
from datetime import (
datetime,
)
import zipfile_zstd
from zipfile import (
ZIP_DEFLATED,
ZIP_STORED,
ZIP_ZSTANDARD,
ZSTANDARD_VERSION,
ZipFile,
ZipInfo,
compressor_names,
)
from pathlib import Path
from taiseilib.common import (
DirPathType,
add_common_args,
run_main,
write_depfile,
)
zstd_decompressor = zstandard.ZstdDecompressor()
def write_zst_file(zf, zst_path, arcname):
'''
Add a file pre-compressed with zstd to the archive
Of course zipfile doesn't support this use-case (because it sucks),
so abuse generous access to its internals to implement it here.
'''
log_file(zst_path, arcname, ZIP_ZSTANDARD)
zip64 = False
with zst_path.open('rb') as zst_file:
zst_size = zst_file.seek(0, 2)
zst_file.seek(0, 0)
zi = ZipInfo.from_file(str(zst_path), arcname=arcname)
zi.compress_type = ZIP_ZSTANDARD
zi.create_version = ZSTANDARD_VERSION
zi.extract_version = ZSTANDARD_VERSION
zi.compress_size = zst_size
if not zi.external_attr:
zi.external_attr = 0o600 << 16 # permissions: ?rw-------
# Unfortunately we must decompress it to compute crc32.
# We'll also compute file size from decompressed data instead of relying on frame headers.
zi.file_size = 0
zi.CRC = 0
for chunk in zstd_decompressor.read_to_iter(zst_file):
zi.file_size += len(chunk)
zi.CRC = zlib.crc32(chunk, zi.CRC)
if zf._seekable:
zf.fp.seek(zf.start_dir)
zi.header_offset = zf.fp.tell()
zf._writecheck(zi)
zf._didModify = True
zf.fp.write(zi.FileHeader(zip64))
zf._writing = True
try:
zst_file.seek(0, 0)
shutil.copyfileobj(zst_file, zf.fp)
assert zst_file.tell() == zi.compress_size
zf.filelist.append(zi)
zf.NameToInfo[zi.filename] = zi
zf.start_dir = zf.fp.tell()
finally:
zf._writing = False
def log_file(path, arcname, comp_type=None):
if str(arcname).endswith('/'):
prefix = 'dir'
else:
prefix = compressor_names.get(comp_type, '???')
print('% 12s' % prefix, '|', arcname, '<--', str(path))
def pack(args):
nocompress_file = args.directory / '.nocompress'
if 1:
comp_type = ZIP_ZSTANDARD
comp_level = 20
else:
comp_type = ZIP_DEFLATED
comp_level = 9
try:
nocompress = list(map(re.compile, filter(None, nocompress_file.read_text().strip().split('\n'))))
except FileNotFoundError:
nocompress = []
nocompress_file = None
zkwargs = {}
if (sys.version_info.major, sys.version_info.minor) >= (3, 7):
zkwargs['compresslevel'] = comp_level
dependencies = []
with ZipFile(str(args.output), 'w', comp_type, **zkwargs) as zf:
for path in sorted(args.directory.glob('**/*')):
if path.name[0] == '.' or any(path.match(x) for x in args.exclude):
continue
relpath = path.relative_to(args.directory)
if path.is_dir():
zi = ZipInfo(str(relpath) + "/", datetime.fromtimestamp(path.stat().st_mtime).timetuple())
zi.compress_type = ZIP_STORED
zi.external_attr = 0o40755 << 16 # drwxr-xr-x
log_file(path, zi.filename)
zf.writestr(zi, '')
else:
dependencies.append(path)
if path.suffix == '.zst':
write_zst_file(zf, path, str(relpath.with_suffix('')))
else:
ctype = comp_type
for pattern in nocompress:
if pattern.match(str(relpath)):
ctype = ZIP_STORED
break
log_file(path, relpath, ctype)
zf.write(str(path), str(relpath), compress_type=ctype)
if args.depfile is not None:
if nocompress_file is not None:
dependencies.append(nocompress_file)
dependencies.append(Path(__file__).resolve())
write_depfile(args.depfile, args.output, dependencies)
def main(args):
import argparse
parser = argparse.ArgumentParser(description='Package game assets.', prog=args[0])
parser.add_argument('directory',
type=DirPathType,
help='the source package directory'
)
parser.add_argument('output',
type=Path,
help='the output archive path'
)
parser.add_argument('--exclude',
action='append',
default=[],
help='file exclusion pattern'
)
add_common_args(parser, depfile=True)
args = parser.parse_args(args[1:])
pack(args)
if __name__ == '__main__':
run_main(main)