499 lines
16 KiB
Python
Executable File
499 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# coding: utf-8
|
||
|
||
import argparse
|
||
import copy
|
||
import hashlib
|
||
import json
|
||
import tarfile
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime
|
||
from pathlib import Path, PurePosixPath
|
||
|
||
|
||
MODALITIES = ("2D", "3D")
|
||
DEFAULT_ARCHIVE_GLOB = "part_*.tar.gz"
|
||
DEFAULT_REPORT_NAME = "rewrite_archives_report.json"
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(
|
||
description=(
|
||
"将已有可视化归档中的成员路径改写为 part_xxxx/images/<filename>，"
|
||
"并输出到新的归档目录。"
|
||
)
|
||
)
|
||
parser.add_argument(
|
||
"source_path",
|
||
help="输入路径，可为 archive root、单个 modality 目录，或单个 tar/tar.gz 文件。",
|
||
)
|
||
parser.add_argument(
|
||
"output_root",
|
||
nargs="?",
|
||
default=None,
|
||
help="输出目录。默认在 source_path 同级生成 <name>_images_prefixed。",
|
||
)
|
||
parser.add_argument(
|
||
"--modalities",
|
||
default="2D,3D",
|
||
help="当 source_path 为 archive root 时要处理的模态，逗号分隔，默认 2D,3D。",
|
||
)
|
||
parser.add_argument(
|
||
"--archive-glob",
|
||
default=DEFAULT_ARCHIVE_GLOB,
|
||
help="归档匹配模式，默认 part_*.tar.gz。",
|
||
)
|
||
parser.add_argument(
|
||
"--parts",
|
||
default=None,
|
||
help=(
|
||
"仅处理指定 part，逗号分隔；既支持 part_0001，也支持 part_0001.tar.gz。"
|
||
),
|
||
)
|
||
parser.add_argument(
|
||
"--limit",
|
||
type=int,
|
||
default=None,
|
||
help="最多处理多少个归档，便于抽样验证。",
|
||
)
|
||
parser.add_argument(
|
||
"--workers",
|
||
type=int,
|
||
default=1,
|
||
help="并发处理归档数，默认 1。",
|
||
)
|
||
parser.add_argument(
|
||
"--gzip-compresslevel",
|
||
type=int,
|
||
default=1,
|
||
help="输出 tar.gz 的 gzip 压缩级别，范围 0-9，默认 1。",
|
||
)
|
||
parser.add_argument(
|
||
"--checksum",
|
||
action="store_true",
|
||
help="输出归档后计算 sha256。",
|
||
)
|
||
parser.add_argument(
|
||
"--verify",
|
||
action="store_true",
|
||
default=True,
|
||
help="输出归档后校验成员数及路径前缀，默认开启。",
|
||
)
|
||
parser.add_argument(
|
||
"--no-verify",
|
||
dest="verify",
|
||
action="store_false",
|
||
help="关闭输出归档校验。",
|
||
)
|
||
parser.add_argument(
|
||
"--overwrite",
|
||
action="store_true",
|
||
help="允许覆盖已有输出归档。",
|
||
)
|
||
parser.add_argument(
|
||
"--report-name",
|
||
default=DEFAULT_REPORT_NAME,
|
||
help=f"输出报告文件名，默认 {DEFAULT_REPORT_NAME}。",
|
||
)
|
||
parser.add_argument(
|
||
"--dry-run",
|
||
action="store_true",
|
||
help="仅输出转换计划，不实际生成归档。",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def now_str():
|
||
return datetime.now().isoformat(timespec="seconds")
|
||
|
||
|
||
def normalize_modalities(raw_modalities):
|
||
modalities = []
|
||
for item in raw_modalities.split(","):
|
||
modality = item.strip()
|
||
if not modality:
|
||
continue
|
||
if modality not in MODALITIES:
|
||
raise ValueError(f"unsupported modality: {modality}")
|
||
modalities.append(modality)
|
||
if not modalities:
|
||
raise ValueError("modalities 不能为空。")
|
||
return tuple(dict.fromkeys(modalities))
|
||
|
||
|
||
def archive_suffix(path):
|
||
name = path.name
|
||
if name.endswith(".tar.gz"):
|
||
return ".tar.gz"
|
||
if name.endswith(".tar"):
|
||
return ".tar"
|
||
raise ValueError(f"unsupported archive suffix: {path}")
|
||
|
||
|
||
def strip_archive_suffix(name):
|
||
if name.endswith(".tar.gz"):
|
||
return name[: -len(".tar.gz")]
|
||
if name.endswith(".tar"):
|
||
return name[: -len(".tar")]
|
||
return name
|
||
|
||
|
||
def parse_parts_filter(raw_parts):
|
||
if not raw_parts:
|
||
return None
|
||
normalized = []
|
||
for item in raw_parts.split(","):
|
||
value = item.strip()
|
||
if not value:
|
||
continue
|
||
normalized.append(strip_archive_suffix(value))
|
||
if not normalized:
|
||
return None
|
||
return set(normalized)
|
||
|
||
|
||
def is_archive_file(path):
|
||
return path.is_file() and path.name.endswith((".tar", ".tar.gz"))
|
||
|
||
|
||
def default_output_root(source_path, resolve_mode):
|
||
source_path = source_path.resolve()
|
||
if resolve_mode == "single_archive":
|
||
base_name = strip_archive_suffix(source_path.name)
|
||
else:
|
||
base_name = source_path.name
|
||
return source_path.parent / f"{base_name}_images_prefixed"
|
||
|
||
|
||
def resolve_archives(source_path, archive_glob, modalities):
|
||
source_path = source_path.resolve()
|
||
|
||
if source_path.is_file():
|
||
if not is_archive_file(source_path):
|
||
raise ValueError(f"source_path is not a supported archive file: {source_path}")
|
||
return "single_archive", [source_path], source_path.parent
|
||
|
||
if not source_path.is_dir():
|
||
raise FileNotFoundError(f"source_path does not exist or is not a directory: {source_path}")
|
||
|
||
direct_archives = sorted(path for path in source_path.glob(archive_glob) if is_archive_file(path))
|
||
if source_path.name in MODALITIES and direct_archives:
|
||
return "modality_dir", direct_archives, source_path.parent
|
||
|
||
archive_paths = []
|
||
for modality in modalities:
|
||
modality_dir = source_path / modality
|
||
if not modality_dir.is_dir():
|
||
continue
|
||
archive_paths.extend(
|
||
sorted(path for path in modality_dir.glob(archive_glob) if is_archive_file(path))
|
||
)
|
||
if archive_paths:
|
||
return "archive_root", archive_paths, source_path
|
||
|
||
raise FileNotFoundError(
|
||
f"no archives found under {source_path} with archive_glob={archive_glob}"
|
||
)
|
||
|
||
|
||
def filter_archives(archive_paths, parts_filter, limit):
|
||
selected = archive_paths
|
||
if parts_filter:
|
||
selected = [
|
||
path for path in selected if strip_archive_suffix(path.name) in parts_filter
|
||
]
|
||
if limit is not None:
|
||
selected = selected[:limit]
|
||
if not selected:
|
||
raise FileNotFoundError("没有匹配到待处理归档，请检查 parts 或 archive_glob。")
|
||
return selected
|
||
|
||
|
||
def build_output_archive_path(archive_path, source_base_dir, output_root):
|
||
return output_root / archive_path.relative_to(source_base_dir)
|
||
|
||
|
||
def build_archive_member_dir(archive_path):
|
||
return f"{strip_archive_suffix(archive_path.name)}/images"
|
||
|
||
|
||
def build_output_member_name(member_name, output_archive):
|
||
return f"{build_archive_member_dir(output_archive)}/{PurePosixPath(member_name).name}"
|
||
|
||
|
||
def open_output_tarfile(output_archive, archive_format, gzip_compresslevel):
|
||
suffix = archive_format or archive_suffix(output_archive)
|
||
if suffix == ".tar.gz":
|
||
return tarfile.open(
|
||
output_archive,
|
||
mode="w:gz",
|
||
compresslevel=gzip_compresslevel,
|
||
)
|
||
return tarfile.open(output_archive, mode="w")
|
||
|
||
|
||
def compute_sha256(file_path):
|
||
digest = hashlib.sha256()
|
||
with open(file_path, "rb") as f:
|
||
while True:
|
||
chunk = f.read(1024 * 1024)
|
||
if not chunk:
|
||
break
|
||
digest.update(chunk)
|
||
return digest.hexdigest()
|
||
|
||
|
||
def verify_rewritten_archive(archive_path, expected_member_count):
|
||
member_count = 0
|
||
samples = []
|
||
expected_prefix = f"{build_archive_member_dir(archive_path)}/"
|
||
with tarfile.open(archive_path, mode="r:*") as tar_obj:
|
||
for member in tar_obj:
|
||
if not member.isfile():
|
||
raise ValueError(
|
||
f"rewritten archive contains non-regular member: {archive_path} -> {member.name}"
|
||
)
|
||
if not member.name.startswith(expected_prefix):
|
||
raise ValueError(
|
||
f"rewritten archive member missing expected prefix {expected_prefix}: "
|
||
f"{archive_path} -> {member.name}"
|
||
)
|
||
member_count += 1
|
||
if len(samples) < 5:
|
||
samples.append(member.name)
|
||
|
||
if member_count != expected_member_count:
|
||
raise ValueError(
|
||
f"rewritten archive member count mismatch: {archive_path}, "
|
||
f"expected={expected_member_count}, actual={member_count}"
|
||
)
|
||
return {
|
||
"member_count": member_count,
|
||
"sample_members": samples,
|
||
}
|
||
|
||
|
||
def rewrite_archive(
|
||
source_archive,
|
||
output_archive,
|
||
gzip_compresslevel,
|
||
verify,
|
||
checksum,
|
||
overwrite,
|
||
dry_run,
|
||
):
|
||
source_archive = source_archive.resolve()
|
||
output_archive = output_archive.resolve()
|
||
if output_archive.exists() and not overwrite:
|
||
raise FileExistsError(f"output archive already exists: {output_archive}")
|
||
|
||
samples_before = []
|
||
samples_after = []
|
||
|
||
if dry_run:
|
||
return {
|
||
"status": "dry_run",
|
||
"source_archive": str(source_archive),
|
||
"output_archive": str(output_archive),
|
||
"member_count": None,
|
||
"sample_members_before": samples_before,
|
||
"sample_members_after": samples_after,
|
||
"output_sha256": None,
|
||
"verification": None,
|
||
}
|
||
|
||
output_archive.parent.mkdir(parents=True, exist_ok=True)
|
||
partial_archive = output_archive.with_name(f"{output_archive.name}.partial")
|
||
if partial_archive.exists():
|
||
partial_archive.unlink()
|
||
|
||
seen_output_names = {}
|
||
member_count = 0
|
||
output_archive_format = archive_suffix(output_archive)
|
||
try:
|
||
with tarfile.open(source_archive, mode="r:*") as input_tar:
|
||
with open_output_tarfile(
|
||
partial_archive,
|
||
output_archive_format,
|
||
gzip_compresslevel,
|
||
) as output_tar:
|
||
for member in input_tar:
|
||
if not member.isfile():
|
||
raise ValueError(
|
||
f"only regular file members are supported: "
|
||
f"{source_archive} -> {member.name}"
|
||
)
|
||
|
||
output_name = build_output_member_name(member.name, output_archive)
|
||
if output_name in seen_output_names:
|
||
raise ValueError(
|
||
"duplicated output member name after rewrite: "
|
||
f"{output_name}, first={seen_output_names[output_name]}, "
|
||
f"duplicate={member.name}"
|
||
)
|
||
seen_output_names[output_name] = member.name
|
||
|
||
if len(samples_before) < 5:
|
||
samples_before.append(member.name)
|
||
if len(samples_after) < 5:
|
||
samples_after.append(output_name)
|
||
|
||
input_file = input_tar.extractfile(member)
|
||
if input_file is None:
|
||
raise ValueError(
|
||
f"failed to extract file object from member: {source_archive} -> {member.name}"
|
||
)
|
||
|
||
output_member = copy.copy(member)
|
||
output_member.name = output_name
|
||
if output_member.pax_headers:
|
||
output_member.pax_headers = dict(output_member.pax_headers)
|
||
|
||
try:
|
||
output_tar.addfile(output_member, input_file)
|
||
finally:
|
||
input_file.close()
|
||
|
||
member_count += 1
|
||
|
||
partial_archive.replace(output_archive)
|
||
finally:
|
||
if partial_archive.exists():
|
||
partial_archive.unlink()
|
||
|
||
verification = verify_rewritten_archive(output_archive, member_count) if verify else None
|
||
output_sha256 = compute_sha256(output_archive) if checksum else None
|
||
return {
|
||
"status": "rewritten",
|
||
"source_archive": str(source_archive),
|
||
"output_archive": str(output_archive),
|
||
"member_count": member_count,
|
||
"sample_members_before": samples_before,
|
||
"sample_members_after": samples_after,
|
||
"output_sha256": output_sha256,
|
||
"verification": verification,
|
||
}
|
||
|
||
|
||
def rewrite_archives(archive_paths, source_base_dir, output_root, args):
|
||
tasks = [
|
||
(archive_path, build_output_archive_path(archive_path, source_base_dir, output_root))
|
||
for archive_path in archive_paths
|
||
]
|
||
if args.workers <= 1 or len(tasks) <= 1:
|
||
return [
|
||
rewrite_archive(
|
||
source_archive,
|
||
output_archive,
|
||
args.gzip_compresslevel,
|
||
args.verify,
|
||
args.checksum,
|
||
args.overwrite,
|
||
args.dry_run,
|
||
)
|
||
for source_archive, output_archive in tasks
|
||
]
|
||
|
||
results = [None] * len(tasks)
|
||
with ThreadPoolExecutor(max_workers=min(args.workers, len(tasks))) as executor:
|
||
future_map = {
|
||
executor.submit(
|
||
rewrite_archive,
|
||
source_archive,
|
||
output_archive,
|
||
args.gzip_compresslevel,
|
||
args.verify,
|
||
args.checksum,
|
||
args.overwrite,
|
||
args.dry_run,
|
||
): index
|
||
for index, (source_archive, output_archive) in enumerate(tasks)
|
||
}
|
||
for future in as_completed(future_map):
|
||
index = future_map[future]
|
||
results[index] = future.result()
|
||
return results
|
||
|
||
|
||
def build_report(args, resolve_mode, source_path, output_root, archive_paths, results):
|
||
return {
|
||
"generated_at": now_str(),
|
||
"source_path": str(Path(source_path).resolve()),
|
||
"output_root": str(Path(output_root).resolve()),
|
||
"resolve_mode": resolve_mode,
|
||
"modalities": list(normalize_modalities(args.modalities)),
|
||
"archive_glob": args.archive_glob,
|
||
"parts": sorted(parse_parts_filter(args.parts) or []),
|
||
"limit": args.limit,
|
||
"workers": args.workers,
|
||
"gzip_compresslevel": args.gzip_compresslevel,
|
||
"checksum": args.checksum,
|
||
"verify": args.verify,
|
||
"overwrite": args.overwrite,
|
||
"dry_run": args.dry_run,
|
||
"selected_archive_count": len(archive_paths),
|
||
"results": results,
|
||
}
|
||
|
||
|
||
def write_json(path, data):
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def main():
|
||
args = parse_args()
|
||
if args.workers <= 0:
|
||
raise ValueError("workers 必须大于 0")
|
||
if not 0 <= args.gzip_compresslevel <= 9:
|
||
raise ValueError("gzip-compresslevel 必须在 0 到 9 之间")
|
||
|
||
source_path = Path(args.source_path).resolve()
|
||
modalities = normalize_modalities(args.modalities)
|
||
parts_filter = parse_parts_filter(args.parts)
|
||
|
||
resolve_mode, archive_paths, source_base_dir = resolve_archives(
|
||
source_path,
|
||
args.archive_glob,
|
||
modalities,
|
||
)
|
||
archive_paths = filter_archives(archive_paths, parts_filter, args.limit)
|
||
|
||
output_root = (
|
||
Path(args.output_root).resolve()
|
||
if args.output_root
|
||
else default_output_root(source_path, resolve_mode)
|
||
)
|
||
|
||
print("")
|
||
print("######################################################################")
|
||
print("# Rewrite visualization archives")
|
||
print("######################################################################")
|
||
print(f"Source path : {source_path}")
|
||
print(f"Output root : {output_root}")
|
||
print(f"Resolve mode : {resolve_mode}")
|
||
print(f"Selected count : {len(archive_paths)}")
|
||
print(f"Workers : {args.workers}")
|
||
print(f"Dry run : {args.dry_run}")
|
||
|
||
results = rewrite_archives(archive_paths, source_base_dir, output_root, args)
|
||
report = build_report(args, resolve_mode, source_path, output_root, archive_paths, results)
|
||
report_path = output_root / args.report_name
|
||
write_json(report_path, report)
|
||
|
||
print(f"Report : {report_path}")
|
||
for item in results:
|
||
print(
|
||
f"[{item['status']}] {item['source_archive']} -> {item['output_archive']}"
|
||
)
|
||
if item["sample_members_before"]:
|
||
print(f" before: {item['sample_members_before'][:3]}")
|
||
print(f" after : {item['sample_members_after'][:3]}")
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|