path: root/grobid_tei_xml/parse.py

import io
import xml.etree.ElementTree as ET
from typing import Any, AnyStr, Dict, List, Optional

from .types import GrobidAddress, GrobidAffiliation, GrobidAuthor, GrobidBiblio, GrobidDocument

xml_ns = "http://www.w3.org/XML/1998/namespace"
ns = "http://www.tei-c.org/ns/1.0"


def _string_to_tree(content: AnyStr) -> ET.ElementTree:
    """
    Helper to consistently parse XML into an ElementTree, whether provided as
    str, bytes, wrapper thereof
    """
    if isinstance(content, str):
        return ET.parse(io.StringIO(content))
    elif isinstance(content, bytes):
        return ET.parse(io.BytesIO(content))
    if isinstance(content, io.StringIO) or isinstance(content, io.BytesIO):
        return ET.parse(content)
    elif isinstance(content, ET.ElementTree):
        return content
    else:
        raise TypeError(f"expected XML as string or bytes, got: {type(content)}")


def _parse_authors(elem: Optional[ET.Element], ns: str = ns) -> List[GrobidAuthor]:
    """
    Internal helper to parse one or more TEI 'author' XML tags into
    GrobidAuthor objects. 'author' could appear in document headers or
    citations.
    """
    if not elem:
        return []

    authors = []
    for author_tag in elem.findall(f".//{{{ns}}}author"):
        persname_tag = author_tag.find(f"./{{{ns}}}persName")
        if persname_tag is None:
            # should we do something else here? it is possible to have author
            # without persName?
            continue

        # basic author name stuff
        given_name = persname_tag.findtext(f"./{{{ns}}}forename") or None
        surname = persname_tag.findtext(f"./{{{ns}}}surname") or None
        # instead create full_name from all the sub-components of the tag
        full_name = " ".join([t.strip() for t in persname_tag.itertext() if t.strip()]).strip()
        ga = GrobidAuthor(
            full_name=full_name or None,
            given_name=given_name,
            surname=surname,
        )

        # author affiliation
        affiliation_tag = author_tag.find(f"./{{{ns}}}affiliation")
        if affiliation_tag is not None:
            affiliation_dict: Dict[str, Any] = dict()
            for orgname_tag in affiliation_tag.findall(f"./{{{ns}}}orgName"):
                orgname_type = orgname_tag.get("type")
                if orgname_type:
                    affiliation_dict[orgname_type] = orgname_tag.text or None
            if affiliation_dict:
                ga.affiliation = GrobidAffiliation(
                    institution=affiliation_dict.get('institution'),
                    department=affiliation_dict.get('department'),
                    laboratory=affiliation_dict.get('laboratory'),
                )
                address_tag = affiliation_tag.find(f"./{{{ns}}}address")
                if address_tag is not None:
                    address_dict = dict()
                    for t in list(address_tag):
                        address_dict[t.tag.split("}")[-1]] = t.text or None
                    if address_dict:
                        ga.affiliation.address = GrobidAddress(
                            addr_line=address_dict.get('addrLine'),
                            post_code=address_dict.get('postCode'),
                            settlement=address_dict.get('settlement'),
                            country=address_dict.get('country'),
                        )
        authors.append(ga)

    return authors


def _clean_url(url: Optional[str]) -> Optional[str]:
    if not url:
        return None
    url = url.strip()
    if url.endswith(".Lastaccessed"):
        url = url.replace(".Lastaccessed", "")
    if url.startswith("<"):
        url = url[1:]
    if ">" in url:
        url = url.split(">")[0]
    return url or None


def test_clean_url() -> None:
    examples: List[dict] = [
        dict(
            dirty="https://archive.org/thing.pdf",
            clean="https://archive.org/thing.pdf",
        ),
        dict(
            dirty="https://archive.org/thing.pdf.Lastaccessed",
            clean="https://archive.org/thing.pdf",
        ),
        dict(
            dirty="<https://archive.org/thing.pdf>",
            clean="https://archive.org/thing.pdf",
        ),
        dict(
            dirty="   https://archive.org/thing.pdf>",
            clean="https://archive.org/thing.pdf",
        ),
        dict(
            dirty="   https://archive.org/thing.pdf>",
            clean="https://archive.org/thing.pdf",
        ),
        dict(dirty="", clean=None),
        dict(dirty=None, clean=None),
    ]

    for row in examples:
        assert row['clean'] == _clean_url(row['dirty'])


def _parse_biblio(elem: ET.Element, ns: str = ns) -> GrobidBiblio:
    """
    Parses an entire TEI 'biblStruct' or 'teiHeader' XML tag

    Could be document header or a citation.
    """

    biblio = GrobidBiblio(
        authors=_parse_authors(elem, ns=ns),
        id=elem.attrib.get("{http://www.w3.org/XML/1998/namespace}id") or None,
        unstructured=elem.findtext(f'.//{{{ns}}}note[@type="raw_reference"]') or None,

        # date below
        title=elem.findtext(f".//{{{ns}}}analytic/{{{ns}}}title") or None,
        journal=elem.findtext(f".//{{{ns}}}monogr/{{{ns}}}title") or None,
        journal_abbrev=None,  # XXX
        publisher=elem.findtext(f".//{{{ns}}}publicationStmt/{{{ns}}}publisher") or None,
        volume=elem.findtext(f'.//{{{ns}}}biblScope[@unit="volume"]') or None,
        issue=elem.findtext(f'.//{{{ns}}}biblScope[@unit="issue"]') or None,
        # pages below
        # XXX: note
        doi=elem.findtext(f'.//{{{ns}}}idno[@type="DOI"]') or None,
        pmid=elem.findtext(f'.//{{{ns}}}idno[@type="PMID"]') or None,
        pmcid=elem.findtext(f'.//{{{ns}}}idno[@type="PMCID"]') or None,
        arxiv_id=elem.findtext(f'.//{{{ns}}}idno[@type="arXiv"]') or None,
        issn=elem.findtext(f'.//{{{ns}}}idno[@type="ISSN"]') or None,
        eissn=elem.findtext(f'.//{{{ns}}}idno[@type="eISSN"]') or None,
    )

    if not biblio.publisher:
        biblio.publisher = elem.findtext(f".//{{{ns}}}imprint/{{{ns}}}publisher") or None

    date_tag = elem.find(f'.//{{{ns}}}date[@type="published"]')
    if date_tag is not None:
        biblio.date = date_tag.attrib.get("when") or None

    # title stuff is messy in references...
    if biblio.journal and not biblio.title:
        biblio.title = biblio.journal
        biblio.journal = None

    if biblio.arxiv_id and biblio.arxiv_id.startswith("arXiv:"):
        biblio.arxiv_id = biblio.arxiv_id[6:]

    el = elem.find(f'.//{{{ns}}}biblScope[@unit="page"]')
    if el is not None:
        if el.attrib.get("from") and el.attrib.get("to"):
            biblio.pages = "{}-{}".format(el.attrib["from"], el.attrib["to"])
        else:
            biblio.pages = el.text

    el = elem.find(f".//{{{ns}}}ptr[@target]")
    if el is not None:
        biblio.url = _clean_url(el.attrib["target"])

    return biblio


def parse_document_xml(xml_text: AnyStr) -> GrobidDocument:
    """
    Use this function to parse TEI-XML of a full document or header processed
    by GROBID.

    Eg, the output of '/api/processFulltextDocument' or '/api/processHeader'
    """
    tree = _string_to_tree(xml_text)
    tei = tree.getroot()

    header = tei.find(f".//{{{ns}}}teiHeader")
    if header is None:
        raise ValueError("XML does not look like TEI format")

    application_tag = header.findall(f".//{{{ns}}}appInfo/{{{ns}}}application")[0]

    doc = GrobidDocument(
        grobid_version=application_tag.attrib["version"].strip(),
        grobid_timestamp=application_tag.attrib["when"].strip(),
        header=_parse_biblio(header),
        pdf_md5=header.findtext(f'.//{{{ns}}}idno[@type="MD5"]') or None,
    )

    refs = []
    for (i, bs) in enumerate(tei.findall(f".//{{{ns}}}listBibl/{{{ns}}}biblStruct")):
        ref = _parse_biblio(bs)
        ref.index = i
        refs.append(ref)
    doc.citations = refs

    text = tei.find(f".//{{{ns}}}text")
    # print(text.attrib)
    if text and text.attrib.get(f"{{{xml_ns}}}lang"):
        doc.language_code = text.attrib[f"{{{xml_ns}}}lang"]  # xml:lang

    el = tei.find(f".//{{{ns}}}profileDesc/{{{ns}}}abstract")
    doc.abstract = (el or None) and " ".join(el.itertext()).strip()
    el = tei.find(f".//{{{ns}}}text/{{{ns}}}body")
    doc.body = (el or None) and " ".join(el.itertext()).strip()
    el = tei.find(f'.//{{{ns}}}back/{{{ns}}}div[@type="acknowledgement"]')
    doc.acknowledgement = (el or None) and " ".join(el.itertext()).strip()
    el = tei.find(f'.//{{{ns}}}back/{{{ns}}}div[@type="annex"]')
    doc.annex = (el or None) and " ".join(el.itertext()).strip()

    return doc


def parse_citations_xml(xml_text: AnyStr) -> List[GrobidBiblio]:
    """
    Use this function to parse TEI-XML of one or more references. This should
    work with either /api/processCitation or /api/processCitationList API
    responses from GROBID

    Note that processed citations are usually returned as a bare XML tag, not a
    full XML document, which means that the TEI xmlns is not set. This requires
    a tweak to all downstream parsing code to handle documents with or without
    the namespace.
    """
    if isinstance(xml_text, bytes):
        xml_text = xml_text.replace(b'xmlns="http://www.tei-c.org/ns/1.0"', b'')
    elif isinstance(xml_text, str):
        xml_text = xml_text.replace('xmlns="http://www.tei-c.org/ns/1.0"', '')
    tree = _string_to_tree(xml_text)
    root = tree.getroot()

    if root.tag == 'biblStruct':
        ref = _parse_biblio(root, ns='')
        ref.index = 0
        return [ref]

    refs = []
    for (i, bs) in enumerate(tree.findall(".//biblStruct")):
        ref = _parse_biblio(bs, ns='')
        ref.index = i
        refs.append(ref)
    return refs