import hashlib import json import os import pathlib import re import struct from argparse import ArgumentParser import numpy as np from PIL import Image from tqdm import tqdm def remove_auxiliary_data(input_file, output_file): with open(input_file, "rb") as f: data = f.read() output_data = bytearray() output_data.extend(data[:8]) # PNG signature # Locate and copy critical chunks start = 8 while start < len(data): length, chunk_type = struct.unpack(">I4s", data[start : start + 8]) start += 8 chunk_data = data[start : start + length] start += length crc = data[start : start + 4] start += 4 if chunk_type in [b"IHDR", b"PLTE", b"IDAT", b"IEND"]: output_data.extend(struct.pack(">I4s", length, chunk_type)) output_data.extend(chunk_data) output_data.extend(crc) with open(output_file, "wb") as f: f.write(output_data) def get_file_id(file_path: str): file_id = file_path.replace("\\", "/") # resource\\global\\{SERVERNAME}\\resource\\template\\xxx.png-> {SERVERNAME}/xxx m = re.search(r"global/(\w+)/resource/template/([^/]+)\.png", file_id) if m: return f"{m.group(1)}/{m.group(2)}" # resource\\template\\xxx.png-> official/xxx m = re.search(r"template/([^/]+)\.png", file_id) if m: return f"official/{m.group(1)}" # resource\\PATH\\TO\\xxx.png-> resource/PATH/TO/xxx m = re.search(r"resource/(.*)\.png", file_id) if m: return f"resource/{m.group(1)}" # docs\\.vuepress\\public\\image\\PATH\\TO\\xxx.png -> docs/PATH/TO/xxx m = re.search(r"docs/.vuepress/public/images/(.*)\.png", file_id) if m: return f"docs/{m.group(1)}" return None def check_png_need_update(file_path: str, perfect_pngs: dict, quiet: bool): # convert to absolute path file_path = str(pathlib.Path(file_path).resolve()) if not file_path.endswith(".png"): return False file_id = get_file_id(file_path) if not file_id: print("unknown file path", file_path) return False if file_id in perfect_pngs: with open(file_path, "rb") as f: data = f.read() sha256 = hashlib.sha256(data).hexdigest() if sha256 == perfect_pngs[file_id]: if not quiet: print("skip", file_path) return False return True def update_png_with_optipng(file_path: str, perfect_pngs: dict, quiet: bool): if check_png_need_update(file_path, perfect_pngs, quiet): if not quiet: print("updating", file_path) else: return False file_path = str(pathlib.Path(file_path).resolve()) file_id = get_file_id(file_path) sz_before = os.stat(file_path).st_size input_file = output_file = file_path remove_auxiliary_data(input_file, output_file) if quiet: os.system(f'optipng -quiet -o7 -zm1-9 -fix "{output_file}"') else: os.system(f'optipng -o7 -zm1-9 -fix "{output_file}"') sz_after = os.stat(file_path).st_size if not quiet: print( f"before: {sz_before} Bytes, after: {sz_after} Bytes, diff: {sz_before - sz_after} Bytes" ) # update file sha256 with open(file_path, "rb") as f: data = f.read() sha256 = hashlib.sha256(data).hexdigest() perfect_pngs.update({file_id: sha256}) if not quiet: print("updated", file_path) update_perfect_png_dict(perfect_pngs) return sz_before - sz_after def update_png_with_oxipng(file_path: str, perfect_pngs: dict, quiet: bool): if check_png_need_update(file_path, perfect_pngs, quiet): if not quiet: print("updating", file_path) else: return False file_path = str(pathlib.Path(file_path).resolve()) file_id = get_file_id(file_path) sz_before = os.stat(file_path).st_size img_before = np.array(Image.open(file_path).convert("L")) if quiet: os.system(f'oxipng -o max --fast -Z -s -q "{file_path}"') os.system(f'oxipng -o 2 -s -q "{file_path}"') else: os.system(f'oxipng -o max --fast -Z -s "{file_path}"') os.system(f'oxipng -o 2 -s "{file_path}"') sz_after = os.stat(file_path).st_size if not quiet: print( f"before: {sz_before} Bytes, after: {sz_after} Bytes, diff: {sz_before - sz_after} Bytes" ) # update file sha256 with open(file_path, "rb") as f: data = f.read() sha256 = hashlib.sha256(data).hexdigest() perfect_pngs.update({file_id: sha256}) img_after = np.array(Image.open(file_path).convert("L")) if not (img_before == img_after).all(): return -1 if not quiet: print("updated", file_path) update_perfect_png_dict(perfect_pngs) return sz_before - sz_after cur_dir = pathlib.Path(__file__).parent.resolve() perfect_pngs_path = str(cur_dir / "optimize_templates.json") def update_perfect_png_dict(perfect_pngs: dict): with open(perfect_pngs_path, "w") as f: json.dump(perfect_pngs, f, indent=4) def ArgParser(): parser = ArgumentParser() parser.add_argument( "-p", "--path", dest="path", nargs="*", help="the paths of png files or folders", default=[], ) parser.add_argument("-q", "--quiet", dest="quiet", action="store_true") return parser if __name__ == "__main__": args = ArgParser().parse_args() paths = args.path quiet = args.quiet perfect_pngs = {} optimized_png_paths = [] if os.path.exists(perfect_pngs_path): with open(perfect_pngs_path, "r") as f: perfect_pngs = json.load(f) if len(paths) == 0: docs_dir = cur_dir.parent.parent / "docs" resource_dir = cur_dir.parent.parent / "resource" paths = [ str(resource_dir / "global"), str(resource_dir / "template"), str(docs_dir), ] print("no path specified, use default paths:", paths) files = [] tqdm_total = 0 for path in paths: if pathlib.Path(path).is_dir(): for root, _, _files in os.walk(path): for f in _files: file = os.path.join(root, f) if check_png_need_update(file, perfect_pngs, quiet): files.append(file) tqdm_total += os.stat(file).st_size ** 2 // 2**20 elif pathlib.Path(path).is_file(): files.append(path) tqdm_total += os.stat(path).st_size ** 2 // 2**20 total_diff_sz = 0 if quiet: t = tqdm(files, total=tqdm_total) for i, file in enumerate(files): cur_file_sz = os.stat(file).st_size diff_sz = update_png_with_oxipng(file, perfect_pngs, quiet) if diff_sz == -1: print("error.") exit(-1) if diff_sz: total_diff_sz += diff_sz optimized_png_paths.append(file) if total_diff_sz < 1024: total_diff_sz_str = f"{total_diff_sz} Bytes" elif total_diff_sz < 1048576: total_diff_sz_str = f"{(total_diff_sz / 1024):.2f} KiB" else: total_diff_sz_str = f"{(total_diff_sz / 1048576):.2f} MiB" if quiet: t.update(cur_file_sz**2 // 2**20) t.set_postfix( file_counts=f"{i+1}/{len(files)}", reduced_size=total_diff_sz_str ) t.refresh() else: print( f"file counts: {i+1}/{len(files)}, reduced pngs size: {total_diff_sz_str}" )