Keep every table in the text, parsed or not, placed or not

src/pubget/_text.py

@@ -3,7 +3,7 @@
 import logging
 import pathlib
 import re
-from typing import Dict, Union
+from typing import Dict, List, Tuple, Union
 
 import pandas as pd
 from lxml import etree
@@ -78,13 +78,106 @@ def extract(
 def _insert_tables(body: str, article_dir: pathlib.Path) -> str:
     """Replace the placeholders left in `body` by the tables' contents.
 
-    Placeholders left for tables that `pubget.extract_articles` did not manage
-    to parse are removed.
+    Every table in the article ends up in the text. One that
+    `pubget.extract_articles` did not manage to parse is rendered from
+    `article.xml` instead, and one the body never places -- a table in
+    `<floats-group>` -- is appended with its caption.
     """
     tables = _load_tables(article_dir)
-    return _TABLE_PLACEHOLDER.sub(
-        lambda match: tables.get(int(match.group(1)), ""), body
-    )
+    from_article = _article_tables(article_dir)
+    placed = set()
+
+    def _substitute(match: "re.Match[str]") -> str:
+        rank = int(match.group(1))
+        if rank in placed:
+            return ""
+        placed.add(rank)
+        return tables.get(rank) or from_article.get(rank, ("", ""))[1]
+
+    body = _TABLE_PLACEHOLDER.sub(_substitute, body)
+    leftover = []
+    for rank in sorted(set(tables) | set(from_article)):
+        if rank in placed:
+            continue
+        caption, rendered = from_article.get(rank, ("", ""))
+        table_text = tables.get(rank) or rendered
+        if table_text:
+            leftover.append("\n".join(filter(None, [caption, table_text])))
+    if not leftover:
+        return body
+    return "\n".join([body.rstrip("\n") + "\n", *leftover])
+
+
+def _article_tables(article_dir: pathlib.Path) -> Dict[int, Tuple[str, str]]:
+    """Each `table-wrap` in `article.xml` as (caption, text), keyed by rank.
+
+    The text has the same shape as `_format_table`'s, read straight from the
+    XML, for the tables `extract_articles` could not parse. The rank counts
+    every `table-wrap` in document order, as the placeholders and the table
+    files both do.
+    """
+    article_file = article_dir.joinpath("article.xml")
+    if not article_file.is_file():
+        return {}
+    try:
+        article = etree.parse(str(article_file))
+    except Exception:
+        _LOG.exception(f"failed to parse {article_file}")
+        return {}
+    tables = {}
+    for rank, wrap in enumerate(article.iter("{*}table-wrap")):
+        label = _text_of(wrap.find("{*}label"))
+        caption = _text_of(wrap.find("{*}caption"))
+        rows = ["\t".join(row) for row in _table_rows(wrap)]
+        foot = _text_of(wrap.find("{*}table-wrap-foot"))
+        parts = [label, "\n".join(rows), foot]
+        table_text = "\n".join(filter(None, parts))
+        tables[rank] = (caption, f"{table_text}\n" if rows else "")
+    return tables
+
+
+def _text_of(elem: "etree._Element | None") -> str:
+    if elem is None:
+        return ""
+    return " ".join("".join(elem.itertext()).split())
+
+
+def _table_rows(wrap: etree._Element) -> List[List[str]]:
+    """The cells of a table, with spanned cells repeated as pandas does."""
+    rows: List[List[str]] = []
+    spans: Dict[int, List] = {}  # column -> [text, rows still covered]
+    for tr in wrap.iter("{*}tr"):
+        row: List[str] = []
+
+        def _fill_spans() -> None:
+            while len(row) in spans:
+                column = len(row)
+                row.append(spans[column][0])
+                spans[column][1] -= 1
+                if spans[column][1] == 0:
+                    del spans[column]
+
+        for cell in tr:
+            if etree.QName(cell).localname not in ("td", "th"):
+                continue
+            _fill_spans()
+            text = _text_of(cell)
+            rowspan = _span(cell, "rowspan")
+            for _ in range(_span(cell, "colspan")):
+                if rowspan > 1:
+                    spans[len(row)] = [text, rowspan - 1]
+                row.append(text)
+        _fill_spans()
+        if any(row):
+            rows.append(row)
+    return rows
+
+
+def _span(cell: etree._Element, attribute: str) -> int:
+    try:
+        return max(1, min(int(cell.get(attribute, 1)), 100))
+    except (TypeError, ValueError):
+        return 1
 
 
 def _load_tables(article_dir: pathlib.Path) -> Dict[int, str]:

tests/test_text.py

@@ -140,25 +140,48 @@ def test_table_text_reproduces_table_csv(tmp_path):
     assert _text._load_tables(tmp_path) == {0: "\tSubject\n\t0012\n"}
 
 
-def test_text_extractor_drops_placeholders_of_missing_tables(
+def test_text_extractor_renders_tables_it_could_not_parse(
     article_with_table,
 ):
-    """Tables that `extract_articles` failed to parse leave no placeholder."""
+    """A table `extract_articles` failed to parse is read from the XML."""
     for table_file in article_with_table.joinpath("tables").glob("table_*"):
         table_file.unlink()
     extractor = _text.TextExtractor(keep_tables=True)
     article = etree.parse(str(article_with_table.joinpath("article.xml")))
     result = extractor.extract(article, article_with_table, {})
-    assert "Peak coordinates." in result["body"]
+    assert _TABLE_TEXT in result["body"]
+    assert result["body"].index("Peak coordinates.") < result["body"].index(
+        _TABLE_TEXT
+    )
     assert "pubget-table" not in result["body"]
 
 
+def test_text_extractor_appends_tables_outside_the_body(tmp_path):
+    """A table in `<floats-group>` has no place in the body; it is appended."""
+    article_dir = tmp_path.joinpath("pmcid_123")
+    article_dir.mkdir()
+    article_dir.joinpath("article.xml").write_bytes(
+        _make_article("<p>Results.</p>").replace(
+            b"</article>",
+            f"<floats-group>{_TABLE}</floats-group></article>".encode(),
+        )
+    )
+    _articles._extract_tables(article_dir)
+    extractor = _text.TextExtractor(keep_tables=True)
+    article = etree.parse(str(article_dir.joinpath("article.xml")))
+    body = extractor.extract(article, article_dir, {})["body"]
+    assert body.index("Results.") < body.index("Peak coordinates.")
+    assert body.index("Peak coordinates.") < body.index(_TABLE_TEXT)
+    assert body.count("IFG\t-42") == 1
+
+
 def test_text_extractor_reports_unreadable_tables(article_with_table, caplog):
     """A table that cannot be read does not prevent extracting the text."""
     article_with_table.joinpath("tables", "table_000.csv").write_text("")
     extractor = _text.TextExtractor(keep_tables=True)
     article = etree.parse(str(article_with_table.joinpath("article.xml")))
     result = extractor.extract(article, article_with_table, {})
     assert "Peak coordinates." in result["body"]
+    assert _TABLE_TEXT in result["body"]
     assert "pubget-table" not in result["body"]
     assert "failed to read table" in caplog.text