rivet/xml_loader.py

from __future__ import annotations

from typing import TYPE_CHECKING

from bs4 import BeautifulSoup

if TYPE_CHECKING:
    from io import TextIOWrapper

class XMLLoader:
    def load(file_: TextIOWrapper) -> dict:
        schema = {}
        bs = BeautifulSoup(file_.read(), "xml")
        structElmts = bs.findAll("structure")
        structures = {}
        
        for structElmt in structElmts:
            structures[structElmt.get("id")] = XMLLoader.parseStructure(structElmt)
        
        schema["structures"] = structures
        return schema
    
    def parseStructure(structElmt: any) -> dict:
        struct = {}
        struct["bits"] = structElmt.get("bits")
        ranges = {}
        rangeElmts = structElmt.findAll("range")
        for rangeElmt in rangeElmts:
            span = rangeElmt.get("end") + "-" + rangeElmt.get("start")
            ranges[span] = XMLLoader.parseRange(rangeElmt)
        
        struct["ranges"] = ranges
        return struct

    def parseRange(rangeElmt: any) -> dict:
        range_ = {}
        range_["name"] = rangeElmt.get("name")
        desc = rangeElmt.find("description")
        if desc is not None: range_["description"] = desc.getText()
        
        valuesElmt = rangeElmt.find("values")
        if valuesElmt is not None:
            range_["values"] = XMLLoader.parseValues(valuesElmt)
        
        if rangeElmt.get("depends-on"):
            range_["depends-on"] = rangeElmt.get("depends-on")
        
        return range_

    def parseValues(valuesElmt: any) -> dict:
        values = {}
        caseElmts = valuesElmt.findAll("case")
        
        for caseElmt in caseElmts:
            val = caseElmt.get("value")
            desc = caseElmt.getText()
            struct = caseElmt.get("structure")
            
            if struct:
                values[val] = {
                    "description": desc,
                    "structure": struct
                }
            else:
                values[val] = desc

        return values