start refactoring into new parser, with dataclass types

8 files changed, 431 insertions, 48 deletions

diff --git a/grobid_tei_xml/__init__.py b/grobid_tei_xml/__init__.py
index 3dc1f76..bf8a133 100644
--- a/grobid_tei_xml/__init__.py
+++ b/grobid_tei_xml/__init__.py
@@ -1 +1,5 @@
 __version__ = "0.1.0"
+
+from .types import GrobidDocument, GrobidCitation
+from .parse import parse_document_xml, parse_citations_xml
+from .grobid2json import teixml2json
diff --git a/grobid_tei_xml/__main__.py b/grobid_tei_xml/__main__.py
new file mode 100644
index 0000000..489bd4e
--- /dev/null
+++ b/grobid_tei_xml/__main__.py
@@ -0,0 +1,29 @@
+
+from .parse import parse_article
+
+def main() -> None:  # pragma no cover
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        description="GROBID TEI XML to JSON",
+        usage="%(prog)s [options] <teifile>...",
+    )
+    parser.add_argument(
+        "--no-encumbered",
+        action="store_true",
+        help=
+        "don't include ambiguously copyright encumbered fields (eg, abstract, body)",
+    )
+    parser.add_argument("teifiles", nargs="+")
+
+    args = parser.parse_args()
+
+    for filename in args.teifiles:
+        content = open(filename, "r").read()
+        print(
+            json.dumps(
+                parse_article(content, encumbered=(not args.no_encumbered)),
+                sort_keys=True,
+            ))
+
+if __name__ == "__main__":  # pragma no cover
+    main()
diff --git a/grobid_tei_xml/grobid2json.py b/grobid_tei_xml/grobid2json.py
index 5edee36..c005b31 100755..100644
--- a/grobid_tei_xml/grobid2json.py
+++ b/grobid_tei_xml/grobid2json.py
@@ -1,5 +1,6 @@
-#!/usr/bin/env python3
 """
+NOTE: this file is DEPRECATED and will be removed soon
+
 NB: adapted to work as a library for PDF extraction. Will probably be
 re-written eventually to be correct, complete, and robust; this is just a
 first iteration.
@@ -24,12 +25,13 @@ Prints JSON to stdout, errors to stderr
 This file copied from the sandcrawler repository.
 """
 
-import argparse
 import io
 import json
 import xml.etree.ElementTree as ET
 from typing import Any, AnyStr, Dict, List, Optional
 
+from .types import *
+
 xml_ns = "http://www.w3.org/XML/1998/namespace"
 ns = "http://www.tei-c.org/ns/1.0"
 
@@ -138,6 +140,7 @@ def biblio_info(elem: ET.Element, ns: str = ns) -> Dict[str, Any]:
     if el is not None:
         ref["url"] = el.attrib["target"]
         # Hand correction
+        # TODO: move this elsewhere
         if ref["url"].endswith(".Lastaccessed"):
             ref["url"] = ref["url"].replace(".Lastaccessed", "")
         if ref["url"].startswith("<"):
@@ -212,31 +215,16 @@ def teixml2json(content: AnyStr, encumbered: bool = True) -> Dict[str, Any]:
             info.pop(k)
     return info
 
-
-def main() -> None:  # pragma no cover
-    parser = argparse.ArgumentParser(
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-        description="GROBID TEI XML to JSON",
-        usage="%(prog)s [options] <teifile>...",
-    )
-    parser.add_argument(
-        "--no-encumbered",
-        action="store_true",
-        help=
-        "don't include ambiguously copyright encumbered fields (eg, abstract, body)",
-    )
-    parser.add_argument("teifiles", nargs="+")
-
-    args = parser.parse_args()
-
-    for filename in args.teifiles:
-        content = open(filename, "r").read()
-        print(
-            json.dumps(
-                teixml2json(content, encumbered=(not args.no_encumbered)),
-                sort_keys=True,
-            ))
-
-
-if __name__ == "__main__":  # pragma no cover
-    main()
+def transform_grobid_ref_xml(raw_xml: str) -> Optional[dict]:
+    """
+    Parses GROBID XML for the case of a single reference/citation string (eg,
+    not a full/propper TEI-XML fulltext document), and returns a dict.
+    """
+    # first, remove any xmlns stuff, for consistent parsing
+    raw_xml = raw_xml.replace('xmlns="http://www.tei-c.org/ns/1.0"', "")
+    tree = ET.parse(io.StringIO(raw_xml))
+    root = tree.getroot()
+    ref = biblio_info(root, ns="")
+    if not any(ref.values()):
+        return None
+    return ref
diff --git a/grobid_tei_xml/grobid_unstructured.py b/grobid_tei_xml/grobid_unstructured.py
index eeb3507..bdead05 100644
--- a/grobid_tei_xml/grobid_unstructured.py
+++ b/grobid_tei_xml/grobid_unstructured.py
@@ -13,19 +13,5 @@ import sys
 import xml.etree.ElementTree as ET
 from typing import Optional
 
-from .grobid2json import biblio_info
+from .parse import biblio_info
 
-
-def transform_grobid_ref_xml(raw_xml: str) -> Optional[dict]:
-    """
-    Parses GROBID XML for the case of a single reference/citation string (eg,
-    not a full/propper TEI-XML fulltext document), and returns a dict.
-    """
-    # first, remove any xmlns stuff, for consistent parsing
-    raw_xml = raw_xml.replace('xmlns="http://www.tei-c.org/ns/1.0"', "")
-    tree = ET.parse(io.StringIO(raw_xml))
-    root = tree.getroot()
-    ref = biblio_info(root, ns="")
-    if not any(ref.values()):
-        return None
-    return ref
diff --git a/grobid_tei_xml/parse.py b/grobid_tei_xml/parse.py
new file mode 100755
index 0000000..a239e4d
--- /dev/null
+++ b/grobid_tei_xml/parse.py
@@ -0,0 +1,204 @@
+
+import io
+import json
+import xml.etree.ElementTree as ET
+from typing import Any, AnyStr, Dict, List, Optional
+
+from .types import *
+
+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:
+    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):
+        return content
+    else:
+        raise TypeError(f"expected XML as string or bytes, got: {type(content)}")
+
+def _parse_authors(elem: Optional[ET.Element]) -> List[GrobidAffiliation]:
+    if not elem:
+        return []
+    names = []
+    for author in elem.findall(f".//{{{ns}}}author"):
+        pn = author.find(f"./{{{ns}}}persName")
+        if not pn:
+            continue
+        given_name = pn.findtext(f"./{{{ns}}}forename") or None
+        surname = pn.findtext(f"./{{{ns}}}surname") or None
+        full_name = " ".join([t.strip() for t in pn.itertext()
+                              if t.strip()]).strip()
+        obj: Dict[str, Any] = dict(name=full_name)
+        if given_name:
+            obj["given_name"] = given_name
+        if surname:
+            obj["surname"] = surname
+        ae = author.find(f"./{{{ns}}}affiliation")
+        if ae:
+            affiliation: Dict[str, Any] = dict()
+            for on in ae.findall(f"./{{{ns}}}orgName"):
+                on_type = on.get("type")
+                if on_type:
+                    affiliation[on_type] = on.text
+            addr_e = ae.find(f"./{{{ns}}}address")
+            if addr_e:
+                address = dict()
+                for t in addr_e.getchildren():
+                    address[t.tag.split("}")[-1]] = t.text
+                if address:
+                    address['post_code'] = address.pop('postCode', None)
+                    affiliation["address"] = GrobidAddress(**address)
+                # previously:
+                # affiliation['address'] = {
+                #    'post_code': addr.findtext('./{%s}postCode' % ns) or None,
+                #    'settlement': addr.findtext('./{%s}settlement' % ns) or None,
+                #    'country': addr.findtext('./{%s}country' % ns) or None,
+                # }
+            obj["affiliation"] = GrobidAffiliation(**affiliation)
+        names.append(GrobidAuthor(**obj))
+    return names
+
+def _parse_citation(elem: ET.Element, ns: str = ns) -> GrobidCitation:
+    ref: Dict[str, Any] = dict()
+    ref["id"] = elem.attrib.get("{http://www.w3.org/XML/1998/namespace}id")
+    ref["unstructured"] = elem.findtext('.//{%s}note[@type="raw_reference"]' %
+                                        ns)
+    # Title stuff is messy in references...
+    ref["title"] = elem.findtext(f".//{{{ns}}}analytic/{{{ns}}}title")
+    other_title = elem.findtext(f".//{{{ns}}}monogr/{{{ns}}}title")
+    if other_title:
+        if ref["title"]:
+            ref["journal"] = other_title
+        else:
+            ref["journal"] = None
+            ref["title"] = other_title
+    ref["authors"] = _parse_authors(elem)
+    ref["publisher"] = elem.findtext(
+        f".//{{{ns}}}publicationStmt/{{{ns}}}publisher")
+    if not ref["publisher"]:
+        ref["publisher"] = elem.findtext(
+            f".//{{{ns}}}imprint/{{{ns}}}publisher")
+    if ref["publisher"] == "":
+        ref["publisher"] = None
+    date = elem.find('.//{%s}date[@type="published"]' % ns)
+    ref["date"] = (date is not None) and date.attrib.get("when")
+    ref["volume"] = elem.findtext('.//{%s}biblScope[@unit="volume"]' % ns)
+    ref["issue"] = elem.findtext('.//{%s}biblScope[@unit="issue"]' % ns)
+    ref["doi"] = elem.findtext('.//{%s}idno[@type="DOI"]' % ns)
+    ref["arxiv_id"] = elem.findtext('.//{%s}idno[@type="arXiv"]' % ns)
+    if ref["arxiv_id"] and ref["arxiv_id"].startswith("arXiv:"):
+        ref["arxiv_id"] = ref["arxiv_id"][6:]
+    ref["pmcid"] = elem.findtext('.//{%s}idno[@type="PMCID"]' % ns)
+    ref["pmid"] = elem.findtext('.//{%s}idno[@type="PMID"]' % ns)
+    el = elem.find('.//{%s}biblScope[@unit="page"]' % ns)
+    if el is not None:
+        if el.attrib.get("from") and el.attrib.get("to"):
+            ref["pages"] = "{}-{}".format(el.attrib["from"], el.attrib["to"])
+        else:
+            ref["pages"] = el.text
+    el = elem.find(".//{%s}ptr[@target]" % ns)
+    if el is not None:
+        ref["url"] = el.attrib["target"]
+        # Hand correction
+        # TODO: move this elsewhere
+        if ref["url"].endswith(".Lastaccessed"):
+            ref["url"] = ref["url"].replace(".Lastaccessed", "")
+        if ref["url"].startswith("<"):
+            ref["url"] = ref["url"][1:]
+        if ">" in ref["url"]:
+            ref["url"] = ref["url"].split(">")[0]
+    else:
+        ref["url"] = None
+    return GrobidCitation(**ref)
+
+def _parse_journal(elem: ET.Element, ns: str = ns) -> GrobidJournal:
+    journal = dict()
+    journal["name"] = elem.findtext(f".//{{{ns}}}monogr/{{{ns}}}title")
+    journal["publisher"] = elem.findtext(
+        f".//{{{ns}}}publicationStmt/{{{ns}}}publisher")
+    if journal["publisher"] == "":
+        journal["publisher"] = None
+    journal["issn"] = elem.findtext('.//{%s}idno[@type="ISSN"]' % ns)
+    journal["eissn"] = elem.findtext('.//{%s}idno[@type="eISSN"]' % ns)
+    journal["volume"] = elem.findtext('.//{%s}biblScope[@unit="volume"]' % ns)
+    journal["issue"] = elem.findtext('.//{%s}biblScope[@unit="issue"]' % ns)
+    journal["abbrev"] = None
+    return GrobidJournal(**journal)
+
+def _parse_header(elem: ET.Element, ns: str = ns) -> GrobidHeader:
+    header = elem
+    info = dict()
+    info["title"] = header.findtext(f".//{{{ns}}}analytic/{{{ns}}}title")
+    info["authors"] = _parse_authors(
+        header.find(f".//{{{ns}}}sourceDesc/{{{ns}}}biblStruct"))
+    info["journal"] = _parse_journal(header)
+    date = header.find(f'.//{{{ns}}}date[@type="published"]')
+    info["date"] = (date is not None) and date.attrib.get("when")
+    info["doi"] = header.findtext(f'.//{{{ns}}}idno[@type="DOI"]')
+    if info["doi"]:
+        info["doi"] = info["doi"].lower()
+    return GrobidHeader(**info)
+
+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()
+    info = dict()
+    encumbered = True
+
+    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_header(header),
+    )
+
+    refs = []
+    for (i, bs) in enumerate(
+            tei.findall(f".//{{{ns}}}listBibl/{{{ns}}}biblStruct")):
+        ref = _parse_citation(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[GrobidCitation]:
+    """
+    Use this function to parse TEI-XML of one or more references.
+    
+    Eg, the output of '/api/processReferences' or '/api/processCitation'.
+    """
+    tree = _string_to_tree(xml_text)
diff --git a/grobid_tei_xml/types.py b/grobid_tei_xml/types.py
new file mode 100644
index 0000000..795d37f
--- /dev/null
+++ b/grobid_tei_xml/types.py
@@ -0,0 +1,96 @@
+
+from typing import Any, AnyStr, Dict, List, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class GrobidAddress:
+    addr_line: Optional[str] = None
+    post_code: Optional[str] = None
+    settlement: Optional[str] = None
+    country: Optional[str] = None
+    country_code: Optional[str] = None
+
+@dataclass
+class GrobidAffiliation:
+    address: Optional[GrobidAddress] = None
+    institution: Optional[str] = None
+    department: Optional[str] = None
+    laboratory: Optional[str] = None
+
+@dataclass
+class GrobidAuthor:
+    name: Optional[str]
+    # TODO: 'forename'?
+    given_name: Optional[str] = None
+    surname: Optional[str] = None
+    affiliation: Optional[dict] = None
+
+@dataclass
+class GrobidCitation:
+    authors: List[GrobidAuthor]
+    index: Optional[int] = None
+    id: Optional[str] = None
+    date: Optional[str] = None
+    issue: Optional[str] = None
+    journal: Optional[str] = None
+    publisher: Optional[str] = None
+    title: Optional[str] = None
+    url: Optional[str] = None
+    volume: Optional[str] = None
+    pages: Optional[str] = None
+    first_page: Optional[str] = None
+    last_page: Optional[str] = None
+    unstructured: Optional[str] = None
+    # TODO: 'arxiv' for consistency?
+    arxiv_id: Optional[str] = None
+    doi: Optional[str] = None
+    pmid: Optional[str] = None
+    pmcid: Optional[str] = None
+    oa_url: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return _simplify_dict(asdict(self))
+
+@dataclass
+class GrobidJournal:
+    name: Optional[str] = None
+    abbrev: Optional[str] = None
+    publisher: Optional[str] = None
+    volume: Optional[str] = None
+    issue: Optional[str] = None
+    issn: Optional[str] = None
+    eissn: Optional[str] = None
+
+@dataclass
+class GrobidHeader:
+    title: Optional[str] = None
+    authors: Optional[str] = None
+    date: Optional[str] = None
+    doi: Optional[str] = None
+    #TODO: note: Optional[str]
+    journal: Optional[GrobidJournal] = None
+
+@dataclass
+class GrobidDocument:
+    grobid_version: str
+    grobid_timestamp: str
+    #TODO: pdf_md5: Optional[str]
+    header: GrobidHeader
+    citations: Optional[List[GrobidCitation]] = None
+    language_code: Optional[str] = None
+    abstract: Optional[str] = None
+    body: Optional[str] = None
+    acknowledgement: Optional[str] = None
+    annex: Optional[str] = None
+
+    def to_dict(self) -> dict:
+        return _simplify_dict(asdict(self))
+
+def _simplify_dict(d: dict) -> dict:
+    for k in list(d.keys()):
+        if isinstance(d[k], dict):
+            d[k] = _simplify_dict(d[k])
+        if d[k] in [None, [], {}, '']:
+            d.pop(k)
+    return d
diff --git a/tests/test_grobid2json.py b/tests/test_grobid2json.py
index 6e3dac2..ed5d996 100644
--- a/tests/test_grobid2json.py
+++ b/tests/test_grobid2json.py
@@ -1,10 +1,12 @@
 import xml
 import json
 import pytest
-from grobid_tei_xml.grobid2json import teixml2json
 
+from grobid_tei_xml import teixml2json, parse_document_xml, GrobidDocument, GrobidCitation
+from grobid_tei_xml.types import *
 
-def test_small_xml():
+
+def test_teixml2json_small_xml():
 
     with open('tests/files/small.xml', 'r') as f:
         tei_xml = f.read()
@@ -13,6 +15,80 @@ def test_small_xml():
 
     assert teixml2json(tei_xml) == json_form
 
+    assert parse_document_xml(tei_xml).to_dict() == json_form
+
+def test_teixml2json_small_xml():
+
+    with open('tests/files/small.xml', 'r') as f:
+        tei_xml = f.read()
+
+    doc = parse_document_xml(tei_xml)
+    expected = GrobidDocument(
+        grobid_version='0.5.1-SNAPSHOT',
+        grobid_timestamp='2018-04-02T00:31+0000',
+        language_code='en',
+        header=GrobidHeader(
+            title="Dummy Example File",
+            authors=[
+                GrobidAuthor(
+                    name="Brewster Kahle",
+                    given_name="Brewster",
+                    surname="Kahle",
+                    affiliation=GrobidAffiliation(
+                        department="Faculty ofAgricultrial Engineering",
+                        laboratory="Plant Physiology Laboratory",
+                        institution="Technion-Israel Institute of Technology",
+                        address=GrobidAddress(
+                            post_code="32000",
+                            settlement="Haifa",
+                            country="Israel",
+                        ),
+                    )
+                ),
+                GrobidAuthor(
+                    name="J Doe",
+                    given_name="J",
+                    surname="Doe",
+                ),
+            ],
+            journal=GrobidJournal(
+                name="Dummy Example File. Journal of Fake News. pp. 1-2. ISSN 1234-5678",
+            ),
+            date="2000",
+        ),
+        abstract="Everything you ever wanted to know about nothing",
+        body="Introduction \nEverything starts somewhere, as somebody [1]  once said. \n\n In Depth \n Meat \nYou know, for kids. \n Potatos \nQED.",
+        citations=[
+            GrobidCitation(
+                index=0,
+                id="b0",
+                authors=[
+                    GrobidAuthor(
+                        name="A Seaperson",
+                        given_name="A",
+                        surname="Seaperson"
+                    )
+                ],
+                date="2001",
+                journal="Letters in the Alphabet",
+                title="Everything is Wonderful",
+                volume="20",
+                pages="1-11",
+            ),
+            GrobidCitation(
+                index=1,
+                id="b1",
+                authors=[],
+                date="2011-03-28",
+                journal="The Dictionary",
+                title="All about Facts",
+                volume="14",
+            ),
+        ],
+    )
+
+    assert doc == expected
+
 
 def test_invalid_xml():
 
diff --git a/tests/test_grobid_unstructured.py b/tests/test_grobid_unstructured.py
index b203b30..91b7398 100644
--- a/tests/test_grobid_unstructured.py
+++ b/tests/test_grobid_unstructured.py
@@ -1,6 +1,6 @@
 import pytest
 
-from grobid_tei_xml.grobid_unstructured import transform_grobid_ref_xml
+from grobid_tei_xml.grobid2json import transform_grobid_ref_xml
 
 
 def test_transform_grobid_ref_xml():