hoodini.utils.classes

  1from __future__ import annotations
  2
  3import xml.etree.ElementTree as ET
  4from pathlib import Path
  5from typing import Any, Literal
  6
  7import polars as pl
  8
  9Mode = Literal["ipg", "nuc2ass", "nucsum", "asssum"]
 10
 11
 12class IPGXMLFile:
 13    def __init__(self, file_path: str | Path):
 14        """Parse an XML file from NCBI IPG/eutils responses."""
 15
 16        self.file_path = Path(file_path)
 17        self.xml_tree = ET.parse(self.file_path)
 18        self.root = self.xml_tree.getroot()
 19
 20    def to_dict(self, mode: Mode) -> dict[str, Any]:
 21        """Return the parsed XML content as a nested dictionary."""
 22
 23        if mode not in ("ipg", "nuc2ass", "nucsum", "asssum"):
 24            raise ValueError(f"Unsupported mode '{mode}'")
 25
 26        parsed_dict: dict[str, Any] = {}
 27
 28        if mode == "ipg":
 29            for ipg_report in self.root.findall("IPGReport"):
 30                ipg_id = ipg_report.get("ipg")
 31                product_acc = ipg_report.get("product_acc")
 32
 33                product = ipg_report.find("Product")
 34                product_details = product.attrib if product is not None else {}
 35
 36                proteins = []
 37                for protein in ipg_report.findall(".//Protein"):
 38                    protein_accver = protein.get("accver")
 39                    protein_details = {k: v for k, v in protein.attrib.items() if k != "accver"}
 40
 41                    cds_list = []
 42                    for cds in protein.findall(".//CDS"):
 43                        cds_list.append(cds.attrib)
 44
 45                    proteins.append(
 46                        {
 47                            "protein_accver": protein_accver,
 48                            "protein_details": protein_details,
 49                            "cds_list": cds_list,
 50                        }
 51                    )
 52
 53                parsed_dict[ipg_id] = {
 54                    "product_acc": product_acc,
 55                    "product_details": product_details,
 56                    "proteins": proteins,
 57                }
 58
 59        elif mode == "nuc2ass":
 60            for link_set in self.root.findall("LinkSet"):
 61                db_from = link_set.findtext("DbFrom")
 62                id_list = [id_tag.text for id_tag in link_set.findall(".//IdList/Id")]
 63
 64                link_set_db = link_set.find("LinkSetDb")
 65                if link_set_db is not None:
 66                    db_to = link_set_db.findtext("DbTo")
 67                    link_name = link_set_db.findtext("LinkName")
 68                    linked_ids = [link.findtext("Id") for link in link_set_db.findall(".//Link")]
 69                else:
 70                    db_to = link_name = None
 71                    linked_ids = []
 72
 73                parsed_dict[tuple(id_list)] = {
 74                    "db_from": db_from,
 75                    "db_to": db_to,
 76                    "link_name": link_name,
 77                    "linked_ids": linked_ids,
 78                }
 79
 80        elif mode == "nucsum":
 81            for doc_sum in self.root.findall("DocSum"):
 82                doc_id = doc_sum.findtext("Id")
 83                doc_items = {item.get("Name"): item.text for item in doc_sum.findall("Item")}
 84                parsed_dict[doc_id] = doc_items
 85
 86        elif mode == "asssum":
 87            for doc_summary in self.root.findall(".//DocumentSummary"):
 88                uid = doc_summary.get("uid")
 89                entry = {child.tag: child.text for child in doc_summary}
 90                parsed_dict[uid] = entry
 91
 92        return parsed_dict
 93
 94    def to_dataframe(self, mode: Mode) -> pl.DataFrame:
 95        """Flatten the parsed XML content into a Polars DataFrame."""
 96
 97        parsed_dict = self.to_dict(mode)
 98        flattened_data: list[dict[str, Any]] = []
 99
100        if mode == "ipg":
101            for ipg_id, details in parsed_dict.items():
102                base_info = {"ipg_id": ipg_id, "product_acc": details["product_acc"]}
103                base_info.update(details["product_details"])
104
105                for protein in details["proteins"]:
106                    protein_info = protein["protein_details"].copy()
107                    protein_info["protein_accver"] = protein["protein_accver"]
108
109                    if protein["cds_list"]:
110                        for cds in protein["cds_list"]:
111                            flattened_data.append({**base_info, **protein_info, **cds})
112                    else:
113                        flattened_data.append({**base_info, **protein_info})
114
115        elif mode == "nuc2ass":
116            for ids, details in parsed_dict.items():
117                id_list_str = ",".join(ids)
118                for linked_id in details["linked_ids"]:
119                    flattened_data.append(
120                        {
121                            "db_from": details["db_from"],
122                            "id_list": id_list_str,
123                            "db_to": details["db_to"],
124                            "link_name": details["link_name"],
125                            "linked_id": linked_id,
126                        }
127                    )
128
129        elif mode == "nucsum":
130            for doc_id, doc_items in parsed_dict.items():
131                row = {"doc_id": doc_id}
132                row.update(doc_items)
133                flattened_data.append(row)
134
135        elif mode == "asssum":
136            for uid, details in parsed_dict.items():
137                row = {"uid": uid}
138                row.update(details)
139                flattened_data.append(row)
140
141        return pl.DataFrame(flattened_data)