311 lines
12 KiB
Python
311 lines
12 KiB
Python
"""danbooru CSV 加工脚本的共享实现。
|
||
|
||
被两个入口脚本复用：
|
||
|
||
- ``modify_danbooru_art.py``：``danbooru_art_full.csv`` -> ``danbooru_art_001.csv``
|
||
- ``modify_danbooru_character.py``：``danbooru_character.csv`` -> ``danbooru_character_001.csv``
|
||
|
||
共同逻辑：按 ``count`` 阈值筛选、只保留指定列、只对指定列做提示词转义、
|
||
临时文件 + 原子替换写回、写回前自动备份。
|
||
|
||
之所以抽成独立模块：**转义规则只有一份实现**。这条规则改过一次（从「处理整行」改成
|
||
「只处理 trigger 列」），两份拷贝迟早会漂移，而它直接决定提示词能不能用。
|
||
|
||
本模块不依赖 ``ra_core`` 与 ComfyUI，可单独导入与测试。
|
||
"""
|
||
|
||
import csv
|
||
import os
|
||
import re
|
||
import shutil
|
||
import time
|
||
from typing import Iterable, List, Optional, Sequence, Tuple
|
||
|
||
__all__ = [
|
||
"BACKUP_SUFFIX",
|
||
"backup_output",
|
||
"backup_path_for",
|
||
"convert",
|
||
"escape_paren",
|
||
"find_columns",
|
||
"normalize_header",
|
||
"read_rows",
|
||
"run",
|
||
"transform_text",
|
||
"write_rows",
|
||
]
|
||
|
||
# 备份后缀：danbooru_art_001.csv -> danbooru_art_001.bak
|
||
BACKUP_SUFFIX = ".bak"
|
||
# 原子替换被占用时的重试次数与间隔（覆盖杀软实时扫描、索引器这类短暂占用）
|
||
REPLACE_RETRIES = 5
|
||
REPLACE_INTERVAL = 0.3
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 文本处理
|
||
# ----------------------------------------------------------------------
|
||
def normalize_header(name: str) -> str:
|
||
"""归一化表头名：去空白/BOM、转小写、空格与短横线转下划线。"""
|
||
return str(name or "").strip().lstrip("\ufeff").lower().replace("-", "_").replace(" ", "_")
|
||
|
||
|
||
def escape_paren(match: "re.Match") -> str:
|
||
"""转义未被转义的括号，已转义的（反斜杠个数为奇数）保持原样。"""
|
||
backslashes, paren = match.group(1), match.group(2)
|
||
if len(backslashes) % 2 == 0:
|
||
return backslashes + "\\" + paren
|
||
return match.group(0)
|
||
|
||
|
||
def transform_text(text: str) -> str:
|
||
"""提示词转义：把源文本整理成可直接使用的标签串。
|
||
|
||
源数据的 trigger 列是**逗号分隔的多标签串**（``hakurei reimu, touhou``），
|
||
因此必须逐段处理，不能对整串做替换——否则逗号后的空格会变成下划线，
|
||
写出 ``hakurei_reimu,_touhou`` 这种脏标签。
|
||
|
||
逐段规则：
|
||
|
||
1. 去首尾空白与短横线（``-x-`` -> ``x``）
|
||
2. 段内连续的空白 / 短横线折叠成**单个**下划线（``95---`` -> ``95``、``a b`` -> ``a_b``）
|
||
3. 折叠重复下划线（``kezune (i- -i)`` -> ``kezune_\\(i_i\\)``）
|
||
4. 空段丢弃（``a,,b`` / 首尾逗号 / 纯空白都收敛掉）
|
||
|
||
最后用 ``,`` 重新连接，并转义未转义的括号（已转义的保持原样）。
|
||
"""
|
||
parts: List[str] = []
|
||
for segment in str(text or "").split(","):
|
||
segment = segment.strip().strip("-").strip()
|
||
segment = re.sub(r"[\s\-]+", "_", segment)
|
||
segment = re.sub(r"_+", "_", segment)
|
||
if segment:
|
||
parts.append(segment)
|
||
return re.sub(r"(\\*)([\(\)])", escape_paren, ",".join(parts))
|
||
|
||
|
||
def find_columns(header: Sequence[str], columns: Sequence[str]) -> dict:
|
||
"""按列名定位各列的索引；缺列时抛出可读异常。"""
|
||
normalized = [normalize_header(cell) for cell in header]
|
||
indexes = {}
|
||
for name in columns:
|
||
if name not in normalized:
|
||
raise ValueError(
|
||
f"表头缺少 {name!r} 列（实际表头：{list(header)}）。"
|
||
f"本工具需要包含 {' / '.join(columns)} 的表头"
|
||
)
|
||
indexes[name] = normalized.index(name)
|
||
return indexes
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 转换
|
||
# ----------------------------------------------------------------------
|
||
def convert(
|
||
rows: Iterable[Sequence[str]], min_count: int, columns: Sequence[str], escape_columns: Sequence[str] = ()
|
||
) -> Tuple[List[List[str]], dict]:
|
||
"""把源行转换为输出行。
|
||
|
||
参数:
|
||
rows: CSV 行迭代器（第一行非空者视为表头）
|
||
min_count: 只保留 ``count`` 不小于该值的行
|
||
columns: 输出的列名（顺序即输出顺序）
|
||
escape_columns: 需要做提示词转义的列名（其余列原样保留）
|
||
|
||
返回:
|
||
``(输出行(含表头), 统计信息)``；统计信息含 ``total`` / ``kept`` / ``skipped``。
|
||
输出按 count 从大到小排序（源文件通常已是降序，这里是防御性排序）。
|
||
"""
|
||
iterator = iter(rows)
|
||
header: Optional[Sequence[str]] = None
|
||
for row in iterator:
|
||
if row and any(str(cell or "").strip() for cell in row):
|
||
header = row
|
||
break
|
||
if header is None:
|
||
raise ValueError("输入文件是空的，找不到表头")
|
||
|
||
indexes = find_columns(header, columns)
|
||
count_index = indexes["count"] if "count" in indexes else None
|
||
if count_index is None:
|
||
raise ValueError(f"输出列里必须包含 count，用于按热度筛选；当前为 {list(columns)}")
|
||
escape_set = {normalize_header(name) for name in escape_columns}
|
||
needs_full_row = max(indexes.values())
|
||
|
||
kept: List[Tuple[int, List[str]]] = []
|
||
stats = {"total": 0, "kept": 0, "skipped": 0}
|
||
|
||
for row in iterator:
|
||
if not row or not any(str(cell or "").strip() for cell in row):
|
||
continue
|
||
stats["total"] += 1
|
||
if needs_full_row >= len(row):
|
||
stats["skipped"] += 1
|
||
continue
|
||
|
||
raw_count = str(row[count_index]).strip()
|
||
if not raw_count.isdigit():
|
||
stats["skipped"] += 1
|
||
continue
|
||
count = int(raw_count)
|
||
if count < min_count:
|
||
continue
|
||
|
||
values: List[str] = []
|
||
empty_critical = False
|
||
for position, name in enumerate(columns):
|
||
key = normalize_header(name)
|
||
if key == "count":
|
||
# count 直接写回规范化后的整数值，不依赖它在第几列
|
||
values.append(str(count))
|
||
continue
|
||
cell = str(row[indexes[name]]).strip()
|
||
if key in escape_set:
|
||
cell = transform_text(cell)
|
||
values.append(cell)
|
||
# 第一列（提示词用的那一列）与需要转义的列都不能为空，否则这行没法使用
|
||
if position == 0 or key in escape_set:
|
||
empty_critical = empty_critical or not cell
|
||
if empty_critical:
|
||
stats["skipped"] += 1
|
||
continue
|
||
|
||
kept.append((count, values))
|
||
|
||
kept.sort(key=lambda item: item[0], reverse=True)
|
||
stats["kept"] = len(kept)
|
||
return [list(columns)] + [row for _, row in kept], stats
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 文件读写
|
||
# ----------------------------------------------------------------------
|
||
def read_rows(path: str) -> List[List[str]]:
|
||
with open(path, encoding="utf-8-sig", newline="") as handle:
|
||
return list(csv.reader(handle))
|
||
|
||
|
||
def write_rows(path: str, rows: Sequence[Sequence[str]]) -> str:
|
||
"""写回文件，返回实际使用的写入方式（供日志说明）。
|
||
|
||
优先「先写临时文件 + 原子替换」，中途失败不会留下半截文件。
|
||
如果目标文件被其他程序占用（编辑器打开、网盘同步、杀软实时扫描），
|
||
先重试几次（覆盖短暂占用）；仍失败则退化为直接覆盖写入——
|
||
此时不具备原子性，但避免了「文件被占用就完全用不了」。
|
||
"""
|
||
temp_path = f"{path}.tmp"
|
||
with open(temp_path, "w", encoding="utf-8", newline="") as handle:
|
||
csv.writer(handle).writerows(rows)
|
||
|
||
last_error: Optional[OSError] = None
|
||
for attempt in range(REPLACE_RETRIES):
|
||
try:
|
||
os.replace(temp_path, path)
|
||
return "原子替换"
|
||
except OSError as error:
|
||
last_error = error
|
||
time.sleep(REPLACE_INTERVAL * (attempt + 1))
|
||
|
||
try:
|
||
with open(path, "w", encoding="utf-8", newline="") as handle:
|
||
csv.writer(handle).writerows(rows)
|
||
except OSError as error:
|
||
# 保留最初原子替换失败的异常（更有诊断价值），并把它链接到当前异常上
|
||
cause = last_error if last_error is not None else error
|
||
raise cause from error
|
||
finally:
|
||
if os.path.exists(temp_path):
|
||
try:
|
||
os.remove(temp_path)
|
||
except OSError:
|
||
pass
|
||
return "直接覆盖（非原子：目标文件被占用）"
|
||
|
||
|
||
def backup_path_for(path: str) -> str:
|
||
"""``danbooru_art_001.csv`` -> ``danbooru_art_001.bak``。"""
|
||
return os.path.splitext(path)[0] + BACKUP_SUFFIX
|
||
|
||
|
||
def backup_output(path: str) -> Optional[str]:
|
||
"""把已有输出文件备份为 ``<输出名>.bak``，返回备份路径。
|
||
|
||
- 输出文件不存在：无需备份，返回 None
|
||
- 备份已存在：不覆盖（保留最初那份原始文件），返回 None
|
||
|
||
本函数不打印任何内容，日志由 :func:`run` 统一输出。
|
||
"""
|
||
if not os.path.isfile(path):
|
||
return None
|
||
target = backup_path_for(path)
|
||
if os.path.exists(target):
|
||
return None
|
||
shutil.copy2(path, target)
|
||
return target
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# 命令行执行
|
||
# ----------------------------------------------------------------------
|
||
def run(
|
||
source: str,
|
||
target: str,
|
||
min_count: int,
|
||
columns: Sequence[str],
|
||
escape_columns: Sequence[str],
|
||
prefix: str,
|
||
backup: bool = True,
|
||
) -> int:
|
||
"""执行一次加工，返回退出码（0 成功）。所有日志都以 ``prefix`` 开头。"""
|
||
source = os.path.abspath(source)
|
||
target = os.path.abspath(target)
|
||
|
||
if min_count < 0:
|
||
print(f"{prefix} --min-count 不能为负数：{min_count}")
|
||
return 1
|
||
if source == target:
|
||
print(f"{prefix} 输入与输出不能是同一个文件：{source}")
|
||
return 1
|
||
if not os.path.isfile(source):
|
||
print(f"{prefix} 输入文件不存在：{source}")
|
||
return 1
|
||
|
||
try:
|
||
rows = read_rows(source)
|
||
output, stats = convert(rows, min_count, columns, escape_columns)
|
||
except UnicodeDecodeError as error:
|
||
print(f"{prefix} 编码无法解析（请使用 UTF-8）：{source} -> {error}")
|
||
return 1
|
||
except (OSError, csv.Error) as error:
|
||
print(f"{prefix} 读取失败：{source} -> {error}")
|
||
return 1
|
||
except ValueError as error:
|
||
print(f"{prefix} {error}")
|
||
return 1
|
||
|
||
try:
|
||
if backup:
|
||
backed = backup_output(target)
|
||
if backed:
|
||
print(f"{prefix} 原文件已备份：{backed}")
|
||
elif os.path.exists(backup_path_for(target)):
|
||
print(f"{prefix} 备份已存在，未覆盖：{backup_path_for(target)}")
|
||
write_mode = write_rows(target, output)
|
||
except OSError as error:
|
||
if isinstance(error, PermissionError):
|
||
print(f"{prefix} 目标文件被其他程序占用，无法写入：{target}")
|
||
print(" 请关闭正在打开该文件的编辑器（或暂停网盘同步 / 杀软实时扫描）后重试；原文件未被修改。")
|
||
print(f"{prefix} 写入失败：{target} -> {error}")
|
||
return 1
|
||
|
||
size = os.path.getsize(target)
|
||
print(f"{prefix} 输入：{source}")
|
||
print(f"{prefix} 输出：{target}（{len(output) - 1} 行 + 表头，{size / 1048576:.2f} MB）")
|
||
print(f"{prefix} 输出列：{', '.join(columns)}；转义列：{', '.join(escape_columns) or '无'}")
|
||
print(
|
||
f"{prefix} 条件：count >= {min_count}"
|
||
f"；源文件 {stats['total']} 行，保留 {stats['kept']} 行，跳过 {stats['skipped']} 行"
|
||
)
|
||
print(f"{prefix} 写入方式：{write_mode}")
|
||
return 0
|