189 lines
4.9 KiB
Python
Executable file
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)
|