ai-cad-quant/experiments/dxf-parser/entity_analysis.py

"""实体分析模块 —— 统计模型空间中各实体类型的数量、INSERT 块聚合、文本提取。"""


def _extract_entity_props(entity):
    """提取单个实体的关键属性为 JSON 友好格式。"""
    dxftype = entity.dxftype()
    props = {}

    if dxftype == "LINE":
        props["start"] = list(entity.dxf.start)
        props["end"] = list(entity.dxf.end)
    elif dxftype == "CIRCLE":
        props["center"] = list(entity.dxf.center)
        props["radius"] = round(entity.dxf.radius, 4)
    elif dxftype == "ARC":
        props["center"] = list(entity.dxf.center)
        props["radius"] = round(entity.dxf.radius, 4)
        props["start_angle"] = round(entity.dxf.start_angle, 1)
        props["end_angle"] = round(entity.dxf.end_angle, 1)
    elif dxftype == "LWPOLYLINE":
        pts = entity.get_points()
        props["vertex_count"] = len(entity)
        props["closed"] = entity.closed
        props["points"] = [[round(p[0], 4), round(p[1], 4)] for p in pts]
    elif dxftype == "POLYLINE":
        verts = list(entity.vertices())
        props["vertex_count"] = len(verts)
        props["points"] = [[round(v.dxf.location[0], 4), round(v.dxf.location[1], 4)] for v in verts]
    elif dxftype == "INSERT":
        props["block_name"] = entity.dxf.name
        props["insert"] = list(entity.dxf.insert)
        props["scale"] = [round(entity.dxf.xscale, 2), round(entity.dxf.yscale, 2), round(entity.dxf.zscale, 2)]
        props["rotation"] = round(entity.dxf.rotation, 2)
    elif dxftype in ("MTEXT", "TEXT"):
        text = entity.plain_text() if dxftype == "MTEXT" else entity.dxf.text
        props["text"] = text
        if hasattr(entity.dxf, "insert"):
            props["insert"] = list(entity.dxf.insert)
        if hasattr(entity.dxf, "height"):
            props["height"] = round(entity.dxf.height, 2)
    elif dxftype == "DIMENSION":
        props["text"] = entity.dxf.text
        if hasattr(entity.dxf, "dimtype"):
            props["dimtype"] = entity.dxf.dimtype
    elif dxftype == "HATCH":
        props["pattern"] = entity.dxf.pattern_name
    elif dxftype in ("SOLID", "TRACE"):
        for vi in range(4):
            try:
                v = getattr(entity.dxf, f"vtx{vi}")
                props[f"vtx{vi}"] = list(v) if v else None
            except Exception:
                break
    elif dxftype == "POINT":
        props["location"] = list(entity.dxf.location)
    elif dxftype == "ELLIPSE":
        props["center"] = list(entity.dxf.center)
        props["major_axis"] = list(entity.dxf.major_axis)
        props["ratio"] = round(entity.dxf.ratio, 4)
    elif dxftype == "SPLINE":
        props["fit_points"] = [list(p) for p in entity.fit_points]
    elif dxftype == "3DFACE":
        for vi in range(4):
            try:
                v = getattr(entity.dxf, f"vtx{vi}")
                props[f"vtx{vi}"] = list(v) if v else None
            except Exception:
                break
    else:
        for attr in entity.dxf.all_existing_dxf_attribs():
            if attr in ("handle", "owner", "layer", "linetype", "color", "lineweight"):
                continue
            val = getattr(entity.dxf, attr)
            if val is not None:
                props[attr] = list(val) if hasattr(val, "__iter__") and not isinstance(val, str) else val

    return props


def analyze_entities(doc) -> dict:
    """分析模型空间中的所有实体。

    Args:
        doc: ezdxf Drawing 对象

    Returns:
        dict: {
            "total_entities": int,
            "type_counts": dict,        # {dxftype: count}
            "blocks": list[dict],        # INSERT 按块名聚合
            "texts": list[dict],         # TEXT/MTEXT 文本列表
        }
    """
    msp = doc.modelspace()
    type_counts = {}

    # INSERT 按块名聚合
    block_map = {}  # block_name -> list of insert_info
    texts = []

    for entity in msp:
        dxftype = entity.dxftype()
        type_counts[dxftype] = type_counts.get(dxftype, 0) + 1

        if dxftype == "INSERT":
            block_name = entity.dxf.name
            insert_info = {
                "handle": entity.dxf.handle,
                "layer": entity.dxf.layer,
                "insert_point": list(entity.dxf.insert),
                "rotation": round(entity.dxf.rotation, 2),
                "scale": [
                    round(entity.dxf.xscale, 2),
                    round(entity.dxf.yscale, 2),
                    round(entity.dxf.zscale, 2),
                ],
            }
            if block_name not in block_map:
                block_map[block_name] = []
            block_map[block_name].append(insert_info)

        elif dxftype in ("TEXT", "MTEXT"):
            content = entity.plain_text() if dxftype == "MTEXT" else entity.dxf.text
            text_info = {
                "handle": entity.dxf.handle,
                "layer": entity.dxf.layer,
                "type": dxftype,
                "content": content,
            }
            if hasattr(entity.dxf, "insert"):
                text_info["insert"] = list(entity.dxf.insert)
            texts.append(text_info)

    # 块聚合结果
    blocks = []
    for block_name, inserts in sorted(block_map.items()):
        blocks.append({
            "block_name": block_name,
            "count": len(inserts),
            "inserts": inserts,
        })

    return {
        "total_entities": sum(type_counts.values()),
        "type_counts": type_counts,
        "blocks": blocks,
        "texts": texts,
    }