Parse UniRef XML -> CSV

parse_uniref_xml.py

#!/usr/bin/python

import os
import sys
import xml
import gzip
import json
import time
from collections import defaultdict
import pandas as pd
from xml.etree.ElementTree import iterparse

fp = sys.argv[1]
output = sys.argv[2]

assert os.path.exists(fp)
assert os.path.exists(output) is False


def parse_member(child):
    member = {}

    for elem in child:
        for attrib in elem:
            member[attrib.attrib["type"].replace(' ', '_')] = attrib.attrib["value"]

    return member


def parse_entry(elem):
    # This particular entry
    entry = defaultdict(set)
    entry["id"].add(elem.get("id"))

    # Members of this entry
    members = []

    for child in elem:
        ending = child.tag.split('}')[-1]

        if ending == "name":
            entry["name"].add(child.text.replace(",", "_"))
        elif ending == "property":
            entry[child.attrib["type"].replace(' ', '_')].add(child.attrib["value"].replace(",", "_"))
        elif ending in ["representativeMember", "member"]:
            for k, v in parse_member(child).items():
                entry[k].add(v.replace(",", "_"))

    entry = {
        k: ';'.join(list(v))
        for k, v in entry.items()
    }

    return entry


def write_csv(
    handle,
    entry=None,
    write_header=False,
    columns=[
        "id",
        "name",
        "common_taxon",
        "common_taxon_ID",
        "UniRef50_ID",
        "GO_Molecular_Function",
        "GO_Biological_Process",
        "GO_Cellular_Component",
        "protein_name",
#        "UniParc_ID",
#        "UniProtKB_accession",
#        "UniRef100_ID"
    ]
):
    if write_header:
        assert entry is None, "Cannot write header and entry together"
        for col in columns:
            assert ',' not in col
        handle.write(",".join(columns) + "\n")
    else:
        assert entry is not None, "Cannot write header and entry together"
        for ix, col in enumerate(columns):
            value = entry.get(col, "")
            assert ',' not in value, "Entry not compatible with CSV format ({})".format(value)
            if ix > 0:
                handle.write(",")
            handle.write(value)
        handle.write("\n")

handle = open(output, "wt")
write_csv(handle, write_header=True)

# Total number of entries
tot_entries = 0

start_time = time.time()
section = None
for event, elem in iterparse(gzip.open(fp), events=["start", "end"]):

    if event == "end" and  elem.tag.split("}")[-1] == "entry":

        entry = parse_entry(elem)
        tot_entries += 1

        write_csv(handle, entry=entry)

        if tot_entries % 100000 == 0:
            print("Processed {:,} entries (time elapsed {:,})".format(
                tot_entries,
                round(time.time() - start_time, 2)))
            start_time = time.time()

        elem.clear()

    elif event == "start" and section is None:
        section = elem.tag
        elem.clear()
    elif event == "end" and elem.tag == section:
        elem.clear()

handle.close()