gdal/nwed.py

from xml.etree.ElementTree import Element
from xml.etree.ElementTree import ElementTree
from xml.etree.ElementTree import SubElement
import attrs
from typing import List, Tuple


@attrs.define
class Nwed:
    # section_name = "002"
    # licheng = "100"
    # remarks_start = "002"
    # pole_name_start = "4001"
    # remarks_end = "003"
    # pole_name_end = "4002"
    # verticalExtent2 = "20"  # 边线宽度
    section_name: str
    licheng: str
    remarks_start: str
    pole_name_start: str
    remarks_end: str
    pole_name_end: str
    verticalExtent2: str  # 边线宽度
    mid_vec_list: List[List[str]]
    right_vec_list: List[List[str]]
    left_vec_list: List[List[str]]

    def indent(self, elem, level=0):
        i = "\n" + level * "\t"
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = i + "\t"
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
            for elem in elem:
                self.indent(elem, level + 1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i

    def AddEleElevationCurve(
        self, elevationCurves, id, vec_list, curve_id, curve_side=None
    ):
        if curve_side:
            elevationCurve_attribs = {"ID": id, "CurveSide": curve_side}
        else:
            elevationCurve_attribs = {"ID": id}
        elevationCurve = SubElement(
            elevationCurves, "EleElevationCurve", attrib=elevationCurve_attribs
        )  # 中线
        SubElement(elevationCurve, "GrowthHeight").text = "0"
        SubElement(elevationCurve, "LinkedSectionObjectID_Group")
        SubElement(elevationCurve, "wallheight").text = "0"
        segmentList = SubElement(elevationCurve, "SegmentList")
        eleElevationCurveSegment = SubElement(
            segmentList,
            "EleElevationCurveSegment",
            # TODO 这个ID其实依据中、边导线会变化。
            attrib={
                "ID": curve_id,
                "StartPointClass": "NotSet",
                "EndPointClass": "NotSet",
            },
        )
        SubElement(eleElevationCurveSegment, "GrowthHeight").text = "0"
        SubElement(eleElevationCurveSegment, "LinkedSectionObjectID_Group")
        vec_list_element = SubElement(eleElevationCurveSegment, "VecList")
        for vec in vec_list:
            SubElement(vec_list_element, "Vector3", attrib={"xyz": ",".join(vec)})
        vec_list_rendering_element = SubElement(
            eleElevationCurveSegment, "VecListRendering"
        )
        for vec in vec_list:
            SubElement(
                vec_list_rendering_element, "Vector3", attrib={"xyz": ",".join(vec)}
            )
        return elevationCurve

    def AddAnnotationElements(self, sectionData, pole_attrib, xyz):
        pdmGraphicsObject = SubElement(
            sectionData, "PdmGraphicsObject", attrib=pole_attrib
        )
        SubElement(pdmGraphicsObject, "GrowthHeight").text = "0"
        SubElement(pdmGraphicsObject, "LinkedSectionObjectID_Group")
        SubElement(pdmGraphicsObject, "Position", attrib={"xyz": xyz})

    def write(self, xml_path):
        root = Element(
            "OverheadTransmissionLineDrawing",
            attrib={
                "Name": self.section_name,
                "xmlns:xsi": "http://www.w3.org/2001/XMLSchema-instance",
                "xmlns:xsd": "http://www.w3.org/2001/XMLSchema",
                "Description": "",
                "Version": "1.0",
                "StartAngle": "0",
                "EndAngle": "0",
                "NextID": "9",
            },
        )
        SubElement(root, "lstCrossingCable")
        drawingFrames = SubElement(root, "DrawingFrames", attrib={"PaperSizeId": "0"})
        SubElement(drawingFrames, "GrowthHeight").text = "0"
        SubElement(drawingFrames, "LinkedSectionObjectID_Group")
        frames = SubElement(drawingFrames, "Frames")
        frame = SubElement(
            frames,
            "Frame",
            attrib={
                "Name": "完整图纸",
                "SplitterPosition": "0",
                "PaperOrientation": "Landscape",
                "UserVerticalOffset": "0",
            },
        )
        SubElement(frame, "GrowthHeight").text = "0"
        SubElement(frame, "LinkedSectionObjectID_Group")
        elevAdjSections = SubElement(frame, "ElevAdjSections")
        SubElement(elevAdjSections, "GrowthHeight").text = "0"
        SubElement(elevAdjSections, "LinkedSectionObjectID_Group")
        sections = SubElement(elevAdjSections, "Sections")
        SubElement(
            sections,
            "ElevationAdjSection",
            attrib={"Position": "0", "ElevationAjusting": "0"},
        )
        SubElement(root, "GPSPointGroups")
        SubElement(root, "Orthoimages")
        SubElement(root, "GPSXianGao")
        SubElement(root, "HouseSwing")
        sectionData = SubElement(root, "SectionData")
        SubElement(sectionData, "licheng").text = self.licheng
        SubElement(sectionData, "lstlicheng")
        SubElement(sectionData, "lstSctnRcd")
        SubElement(sectionData, "ForestElevationCurves")
        elevationCurves = SubElement(sectionData, "ElevationCurves")
        self.AddEleElevationCurve(elevationCurves, "1", self.mid_vec_list, "4")
        self.AddEleElevationCurve(elevationCurves, "2", self.left_vec_list, "5", "Left")
        self.AddEleElevationCurve(
            elevationCurves, "3", self.right_vec_list, "6", "Right"
        )
        annotationElements = SubElement(sectionData, "AnnotationElements")
        self.AddAnnotationElements(
            annotationElements,
            {
                "xsi:type": "ElePowerLinePole",
                "ID": "7",
                "CurveSide": "All",
                "Remarks": self.remarks_start,
                "PoleType": "AnglePole",
                "LeftElevation": "0",
                "RightElevation": "0",
                "CableAngle": "0",
                "PoleName": self.pole_name_start,
            },
            xyz=",".join(self.mid_vec_list[0]),
        )
        self.AddAnnotationElements(
            annotationElements,
            {
                "xsi:type": "ElePowerLinePole",
                "ID": "8",
                "CurveSide": "All",
                "Remarks": self.remarks_end,
                "PoleType": "AnglePole",
                "LeftElevation": "0",
                "RightElevation": "0",
                "CableAngle": "0",
                "PoleName": self.pole_name_end,
            },
            xyz=",".join(self.mid_vec_list[-1]),
        )
        mapData = SubElement(root, "MapData")
        SubElement(mapData, "GrowthHeight").text = "0"
        SubElement(mapData, "LinkedSectionObjectID_Group")
        SubElement(mapData, "verticalExtent2").text = self.verticalExtent2
        SubElement(mapData, "MapElements")
        tree = ElementTree(root)
        self.indent(root)
        tree.write(xml_path, encoding="utf-8", xml_declaration=True)