gwdoc/resetlevel.py

import copy
import json
import re
import sys
from typing import Any


CN_NUM = "\u4e00\u4e8c\u4e09\u56db\u4e94\u516d\u4e03\u516b\u4e5d\u5341\u767e\u5343\u4e07\u3007\u96f6\u4e24"

CHAPTER_PATTERN = re.compile(rf"^\s*\u7b2c\s*([{CN_NUM}\d]+)\s*\u7ae0")
CN_LEVEL1_PATTERN = re.compile(rf"^\s*([{CN_NUM}]+)(?=\s|[\u3001\u3002\uff0e.])")

CASE_LIKE_PATTERN = re.compile(r"^\s*\d+(?:\.|\uff0e)(?!\d|\s|$)")
ARABIC_PATTERN = re.compile(
    r"^\s*(\d+(?:\.\d+)*)(?=\s|[\u4e00-\u9fff\u3001\u3002\uff0e]|$|\.(?:\s|$))"
)


def has_valid_text_level(item: dict) -> bool:
    return "text_level" in item and item.get("text_level") not in (None, "")


def parse_heading(text: str):
    text = text.strip()
    if not text:
        return None, None

    if CASE_LIKE_PATTERN.match(text):
        return "case_like", None

    if CHAPTER_PATTERN.match(text):
        return "level1", None

    if CN_LEVEL1_PATTERN.match(text):
        return "level1", None

    match = ARABIC_PATTERN.match(text)
    if match:
        return "arabic", match.group(1)

    return None, None


def arabic_depth(number: str) -> int:
    return number.count(".") + 1


def parent_number(number: str) -> str | None:
    parts = number.split(".")
    if len(parts) <= 1:
        return None
    return ".".join(parts[:-1])


def collect_headings(data: list) -> list:
    headings = []

    for item in data:
        if not isinstance(item, dict):
            continue
        if item.get("type") != "text":
            continue
        if not has_valid_text_level(item):
            continue

        kind, number = parse_heading(item.get("text", ""))
        if kind in {"level1", "arabic", "case_like"}:
            headings.append({"kind": kind, "number": number})

    return headings


def is_valid_arabic_hierarchy(numbers: list[str], allow_missing_depth1_parent: bool) -> bool:
    seen = set()
    has_depth2 = False
    has_depth3 = False

    for number in numbers:
        depth = arabic_depth(number)
        if depth >= 2:
            has_depth2 = True
        if depth >= 3:
            has_depth3 = True

        parent = parent_number(number)
        if parent is not None:
            parent_depth = arabic_depth(parent)
            parent_is_missing_depth1 = allow_missing_depth1_parent and parent_depth == 1
            if parent not in seen and not parent_is_missing_depth1:
                return False

        seen.add(number)

    return has_depth2 and has_depth3


def detect_valid_mode(headings: list):
    if not headings:
        return None

    if any(heading["kind"] == "case_like" for heading in headings):
        return None

    first = headings[0]

    if first["kind"] == "arabic":
        numbers = [heading["number"] for heading in headings if heading["kind"] == "arabic"]
        if numbers and arabic_depth(numbers[0]) == 1:
            if is_valid_arabic_hierarchy(numbers, allow_missing_depth1_parent=False):
                return "A"
        return None

    if first["kind"] == "level1":
        numbers = [heading["number"] for heading in headings[1:] if heading["kind"] == "arabic"]
        if not numbers:
            return None

        first_depth = arabic_depth(numbers[0])
        if first_depth == 2 and is_valid_arabic_hierarchy(numbers, allow_missing_depth1_parent=True):
            return "B"

        if first_depth == 1 and is_valid_arabic_hierarchy(numbers, allow_missing_depth1_parent=False):
            return "C"

    return None


JsonContent = Any


def reset_textlevel(json_content: JsonContent) -> JsonContent:
    """
    输入 JSON 内容，返回更新 text_level 后的新 JSON 内容。

    json_content 可以是已经 json.load/json.loads 后的 Python 对象，
    也可以是 JSON 字符串。函数不会修改原始入参。
    """
    if isinstance(json_content, str):
        data = json.loads(json_content)
    else:
        data = copy.deepcopy(json_content)

    if not isinstance(data, list):
        return data

    headings = collect_headings(data)
    mode = detect_valid_mode(headings)

    if mode is None:
        return data

    for item in data:
        if not isinstance(item, dict):
            continue
        if item.get("type") != "text":
            continue
        if not has_valid_text_level(item):
            continue

        kind, number = parse_heading(item.get("text", ""))
        if kind == "level1":
            item["text_level"] = 1
            continue
        if kind != "arabic":
            continue

        depth = arabic_depth(number)
        item["text_level"] = depth if mode in {"A", "B"} else depth + 1

    return data





if __name__ == "__main__":
    json_file_path = r'E:\ZKYNLP\Hjunproject\project0506\kgrag\163-06A0014-B01001_发动机-维修手册_content_list.json'
    with open(json_file_path, "r", encoding="utf-8") as f:
        json_content = json.load(f)
    result=  reset_textlevel(json_content)

    for ins in result[:30]:
        print(ins)
    # print(json.dumps(result, ensure_ascii=False, indent=2))