[ENH] make cross references configurable

README.md

@@ -253,6 +253,10 @@ If we had not used `--articles_with_coords_only`, the new subdirectory would be
   Fields are `pmcid`, `surname`, `given-names`.
 - `text.csv` contains one row per article.
   The first field is the `pmcid`, and the other fields are `title`, `keywords`, `abstract`, and `body`, and contain the text extracted from these parts of the article.
+  By default the tables are left out of the `body` (they are available in the CSV files created by `pubget extract_articles`, described above).
+  If we use the `--keep_tables` option, each table is inserted in the `body` at the position where it appears in the article: its label, then its contents as tab-separated values (with one line per header row), then its footer.
+  The cells are those of the table's CSV file, so the text and the CSV files always agree.
+  Tables that `pubget` failed to parse are still left out.
 - `links.csv` contains the external links found in the articles.
   The fields are `pmcid`, `ext-link-type` (the type of link, for example "uri", "doi"), and `href` (usually an URL).
 - `neurovault_collections.csv` and `neurovault_images.csv`: [NeuroVault](https://neurovault.org/) collection and image IDs that could be extracted from links in the articles, if any.

src/pubget/_data/stylesheets/text_extraction.xsl

@@ -7,6 +7,8 @@
 
   <xsl:output method="xml" version="1.0" encoding="UTF-8" omit-xml-declaration="no"/>
   <xsl:strip-space elements="*"/>
+  <xsl:param name="preserve-crossrefs" select="'true'"/>
+  <xsl:param name="keep-tables" select="'false'"/>
 
   <xsl:template match="/">
     <extracted-text>
@@ -78,6 +80,22 @@
   </xsl:text>
 </xsl:template>
 
+<!-- The table itself is always stripped from the text; when tables are kept a
+     placeholder is left in its place, to be replaced with the table's contents
+     by 'pubget._text'. The number is the table's rank in the article, which is
+     also the number used to name the files written for it by
+     'pubget.extract_articles'. -->
+<xsl:template match="table-wrap" >
+  <xsl:text> </xsl:text>
+  <xsl:apply-templates />
+  <xsl:if test="$keep-tables = 'true'">
+    <xsl:text>&#10;[pubget-table-</xsl:text>
+    <xsl:value-of select="count(preceding::table-wrap)"/>
+    <xsl:text>]&#10;</xsl:text>
+  </xsl:if>
+  <xsl:text> </xsl:text>
+</xsl:template>
+
 <xsl:template match="text()" >
   <xsl:copy-of select="."/>
   <xsl:text> </xsl:text>
@@ -305,6 +323,13 @@
   <xsl:template match="volume-series" />
   <xsl:template match="word-count" />
   <xsl:template match="x" />
-  <xsl:template match="xref" />
+  <xsl:template match="xref">
+    <xsl:choose>
+      <xsl:when test="$preserve-crossrefs = 'true'">
+        <xsl:apply-templates />
+      </xsl:when>
+      <xsl:otherwise />
+    </xsl:choose>
+  </xsl:template>
   <xsl:template match="year" />
 </xsl:transform>

src/pubget/_data_extraction.py

@@ -139,6 +139,8 @@ def extract_data_to_csv(
     output_dir: Optional[PathLikeOrStr] = None,
     *,
     articles_with_coords_only: bool = False,
+    preserve_cross_references: bool = True,
+    keep_tables: bool = False,
     n_jobs: int = 1,
 ) -> Tuple[Path, ExitCode]:
     """Extract text and coordinates from articles and store in csv files.
@@ -158,6 +160,15 @@ def extract_data_to_csv(
         `articles_with_coords_only`.
     articles_with_coords_only
         If True, articles that contain no stereotactic coordinates are ignored.
+    preserve_cross_references
+        If True, text from inline cross-reference elements is preserved in
+        extracted text. If False, those elements are removed.
+    keep_tables
+        If True, the articles' tables are inserted in the extracted text, at
+        the position where they appear in the article: their label, then their
+        contents as tab-separated values, then their footer. If False, tables
+        are only available in the separate CSV files created by
+        `pubget.extract_articles`.
     n_jobs
         Number of processes to run in parallel. `-1` means using all
         processors.
@@ -190,7 +201,12 @@ def extract_data_to_csv(
     )
     n_jobs = _utils.check_n_jobs(n_jobs)
     n_articles = _do_extract_data_to_csv(
-        articles_dir, output_dir, articles_with_coords_only, n_jobs=n_jobs
+        articles_dir,
+        output_dir,
+        articles_with_coords_only,
+        preserve_cross_references=preserve_cross_references,
+        keep_tables=keep_tables,
+        n_jobs=n_jobs,
     )
     is_complete = bool(status["previous_step_complete"])
     _utils.write_info(
@@ -204,11 +220,17 @@ def extract_data_to_csv(
     return output_dir, exit_code
 
 
-def _get_data_extractors() -> List[Extractor]:
+def _get_data_extractors(
+    preserve_cross_references: bool,
+    keep_tables: bool,
+) -> List[Extractor]:
     return [
         MetadataExtractor(),
         AuthorsExtractor(),
-        TextExtractor(),
+        TextExtractor(
+            preserve_cross_references=preserve_cross_references,
+            keep_tables=keep_tables,
+        ),
         TableInfoExtractor(),
         CoordinateExtractor(),
         CoordinateSpaceExtractor(),
@@ -221,14 +243,19 @@ def _do_extract_data_to_csv(
     articles_dir: Path,
     output_dir: Path,
     articles_with_coords_only: bool,
+    *,
+    preserve_cross_references: bool,
+    keep_tables: bool,
     n_jobs: int,
 ) -> int:
     """Do the data extraction and return the number of articles whose data was
     saved. If `articles_with_coords_only` only articles with at least one
     sterotactic coordinate triplet have their data saved.
     """
     n_to_process = _utils.get_n_articles(articles_dir)
-    data_extractors = _get_data_extractors()
+    data_extractors = _get_data_extractors(
+        preserve_cross_references, keep_tables
+    )
     all_writers = [
         CSVWriter.from_extractor(extractor, output_dir)
         for extractor in data_extractors
@@ -289,6 +316,18 @@ def _edit_argument_parser(
         help="Only keep data for articles in which stereotactic coordinates "
         "are found.",
     )
+    argument_parser.add_argument(
+        "--strip-cross-references",
+        action="store_true",
+        help="Remove inline cross-reference text from extracted article text.",
+    )
+    argument_parser.add_argument(
+        "--keep_tables",
+        action="store_true",
+        help="Insert the articles' tables in the extracted text, at the "
+        "position where they appear in the article: their label, then their "
+        "contents as tab-separated values, then their footer.",
+    )
     _utils.add_n_jobs_argument(argument_parser)
 
 
@@ -309,6 +348,8 @@ def run(
         output_dir, exit_code = extract_data_to_csv(
             previous_steps_output["extract_articles"],
             articles_with_coords_only=args.articles_with_coords_only,
+            preserve_cross_references=not args.strip_cross_references,
+            keep_tables=args.keep_tables,
             n_jobs=args.n_jobs,
         )
         if not _utils.get_n_articles(output_dir):
@@ -343,5 +384,7 @@ def run(self, args: argparse.Namespace) -> ExitCode:
         return extract_data_to_csv(
             args.articles_dir,
             articles_with_coords_only=args.articles_with_coords_only,
+            preserve_cross_references=not args.strip_cross_references,
+            keep_tables=args.keep_tables,
             n_jobs=args.n_jobs,
         )[1]

src/pubget/_text.py

@@ -1,36 +1,66 @@
 """Extracting text from XML articles."""
+import json
 import logging
 import pathlib
+import re
 from typing import Dict, Union
 
+import pandas as pd
 from lxml import etree
 
 from pubget import _utils
 from pubget._typing import Extractor, Records
 
 _LOG = logging.getLogger(__name__)
 
+# Left in the text by 'text_extraction.xsl' at the position of each table when
+# tables are kept, and by the files written for each table by
+# `pubget.extract_articles`. In both cases the number is the table's rank in
+# the article, which is how the two are matched.
+_TABLE_PLACEHOLDER = re.compile(r"\[pubget-table-(\d+)\]")
+_TABLE_INFO_FILE = re.compile(r"table_(\d+)_info\.json")
+# The name pandas gives to a column whose header cell is empty.
+_UNNAMED_COLUMN = re.compile(r"^Unnamed: \d+(_level_\d+)?$")
+
 
 class TextExtractor(Extractor):
     """Extracting text from XML articles."""
 
     fields = ("pmcid", "title", "keywords", "abstract", "body")
     name = "text"
 
+    def __init__(
+        self,
+        preserve_cross_references: bool = True,
+        keep_tables: bool = False,
+    ) -> None:
+        self.preserve_cross_references = preserve_cross_references
+        self.keep_tables = keep_tables
+
     def extract(
         self,
         article: etree.ElementTree,
         article_dir: pathlib.Path,
         previous_extractors_output: Dict[str, Records],
     ) -> Dict[str, Union[str, int]]:
-        del article_dir, previous_extractors_output
+        del previous_extractors_output
         result: Dict[str, Union[str, int]] = {}
         # Stylesheet is not parsed in init because lxml.XSLT cannot be pickled
         # so that would prevent the extractor from being passed to
         # multiprocessing map. Parsing is cached.
         stylesheet = _utils.load_stylesheet("text_extraction.xsl")
         try:
-            transformed = stylesheet(article)
+            transformed = stylesheet(
+                article,
+                **{
+                    "preserve-crossrefs": etree.XSLT.strparam(
+                        "true" if self.preserve_cross_references else "false"
+                    ),
+                    "keep-tables": etree.XSLT.strparam(
+                        "true" if self.keep_tables else "false"
+                    ),
+                },
+            )
         except Exception:
             _LOG.exception(
                 f"failed to transform article: {stylesheet.error_log}"
@@ -39,5 +69,80 @@ def extract(
         for part_name in self.fields:
             elem = transformed.find(part_name)
             result[part_name] = elem.text
+        if self.keep_tables and result["body"]:
+            result["body"] = _insert_tables(str(result["body"]), article_dir)
         result["pmcid"] = _utils.get_pmcid(article)
         return result
+
+
+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.
+    """
+    tables = _load_tables(article_dir)
+    return _TABLE_PLACEHOLDER.sub(
+        lambda match: tables.get(int(match.group(1)), ""), body
+    )
+
+
+def _load_tables(article_dir: pathlib.Path) -> Dict[int, str]:
+    """Read the tables extracted from an article by `extract_articles`.
+
+    Keys are the tables' rank in the article. Tables that cannot be read are
+    left out.
+    """
+    tables = {}
+    for info_file in _utils.get_table_info_files_from_article_dir(article_dir):
+        match = _TABLE_INFO_FILE.match(info_file.name)
+        assert match is not None
+        try:
+            tables[int(match.group(1))] = _format_table(info_file)
+        except Exception:
+            _LOG.exception(f"failed to read table {info_file}")
+    return tables
+
+
+def _format_table(info_file: pathlib.Path) -> str:
+    """Render one of an article's extracted tables as tab-separated values.
+
+    The label and the footer are added because, unlike the caption, they are
+    not part of the extracted text.
+    """
+    table_info = json.loads(info_file.read_text("UTF-8"))
+    table_data = _read_table_data(
+        info_file.with_name(table_info["table_data_file"]),
+        table_info["n_header_rows"],
+    )
+    grid: str = table_data.to_csv(
+        sep="\t", index=False, header=False, lineterminator="\n"
+    )
+    parts = [
+        table_info["table_label"],
+        grid.strip("\n"),
+        table_info["table_foot"],
+    ]
+    table_text = "\n".join(filter(None, parts))
+    return f"{table_text}\n"
+
+
+def _read_table_data(
+    table_csv: pathlib.Path, n_header_rows: int
+) -> pd.DataFrame:
+    """Read a table's CSV file, header rows included, as text.
+
+    Unlike `_utils.read_article_table`, this does not attempt to convert the
+    cells: they are reproduced exactly as `extract_articles` wrote them,
+    rather than passed through pandas' type inference a second time. The
+    header rows are read as regular rows so they keep their original position
+    and the placeholder names pandas gives to unnamed columns can be removed.
+    """
+    table_data = pd.read_csv(
+        table_csv, header=None, dtype=str, keep_default_na=False
+    )
+    header = table_data.iloc[:n_header_rows]
+    table_data.iloc[:n_header_rows] = header.replace(
+        _UNNAMED_COLUMN, "", regex=True
+    )
+    return table_data

tests/test_data_extraction.py

@@ -99,6 +99,16 @@ def test_extract_data_to_csv(
     _check_extracted_data(data_dir, articles_with_coords_only)
 
 
+def test_extract_data_to_csv_with_tables(tmp_path, articles_dir):
+    data_dir, code = _data_extraction.extract_data_to_csv(
+        articles_dir, tmp_path.joinpath("extracted_data"), keep_tables=True
+    )
+    assert code == ExitCode.COMPLETED
+    text = pd.read_csv(data_dir.joinpath("text.csv"))
+    assert text.at[0, "body"].strip().startswith("The text")
+    assert "X\tY\tZ\n10\t20\t30\n" in text.at[0, "body"]
+
+
 def test_extractor_failures(articles_dir, tmp_path, monkeypatch):
     data_dir = Path(f"{tmp_path}-extraction_failures-extracted_data")
     mock = Mock(side_effect=ValueError)
@@ -135,7 +145,12 @@ def test_should_write(data, with_coords, expected):
 
 def test_stop_pipeline(empty_articles_dir):
     step = _data_extraction.DataExtractionStep()
-    args = argparse.Namespace(articles_with_coords_only=False, n_jobs=1)
+    args = argparse.Namespace(
+        articles_with_coords_only=False,
+        strip_cross_references=False,
+        keep_tables=False,
+        n_jobs=1,
+    )
     previous_steps = {"extract_articles": empty_articles_dir}
     with pytest.raises(_typing.StopPipeline, match=r"No articles.*"):
         step.run(args, previous_steps)