diff --git a/README.md b/README.md index 29d134d..2d2dbca 100644 --- a/README.md +++ b/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. diff --git a/src/pubget/_data/stylesheets/text_extraction.xsl b/src/pubget/_data/stylesheets/text_extraction.xsl index 885254c..82e874e 100644 --- a/src/pubget/_data/stylesheets/text_extraction.xsl +++ b/src/pubget/_data/stylesheets/text_extraction.xsl @@ -7,6 +7,8 @@ + + @@ -78,6 +80,22 @@ + + + + + + [pubget-table- + + ] + + + + @@ -305,6 +323,13 @@ - + + + + + + + + diff --git a/src/pubget/_data_extraction.py b/src/pubget/_data_extraction.py index 0a6911d..dc9ebe3 100644 --- a/src/pubget/_data_extraction.py +++ b/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,6 +243,9 @@ 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 @@ -228,7 +253,9 @@ def _do_extract_data_to_csv( 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] diff --git a/src/pubget/_text.py b/src/pubget/_text.py index 847544d..9850b52 100644 --- a/src/pubget/_text.py +++ b/src/pubget/_text.py @@ -1,8 +1,11 @@ """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 @@ -10,6 +13,15 @@ _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.""" @@ -17,20 +29,38 @@ class TextExtractor(Extractor): 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 diff --git a/tests/test_data_extraction.py b/tests/test_data_extraction.py index 94734e5..536a2ac 100644 --- a/tests/test_data_extraction.py +++ b/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) diff --git a/tests/test_text.py b/tests/test_text.py index 19a30ed..dfdcefd 100644 --- a/tests/test_text.py +++ b/tests/test_text.py @@ -1,8 +1,67 @@ +import json +from pathlib import Path from unittest.mock import Mock +import pytest from lxml import etree -from pubget import _text +from pubget import _articles, _text, _utils + +_TABLE = """ + +

Peak coordinates.

+ + + + + + + + + +
MNI
Regionxy
IFG-4218
-4020
+ +

IFG = inferior frontal gyrus.

 + +""" + +_TABLE_TEXT = """Table 1 +\tMNI\tMNI +Region\tx\ty +IFG\t-42\t18 +IFG\t-40\t20 +IFG = inferior frontal gyrus. +""" + + +def _make_article(body: str) -> bytes: + article = f"""
+ + 123 + + {body} +
""" + return article.encode("UTF-8") + + +@pytest.fixture(autouse=True) +def clear_stylesheet_cache(): + """Avoid leaking mocked stylesheets between tests.""" + _utils.load_stylesheet.cache_clear() + yield + _utils.load_stylesheet.cache_clear() + + +@pytest.fixture +def article_with_table(tmp_path): + """An article dir with tables, as created by `extract_articles`.""" + article_dir = tmp_path.joinpath("pmcid_123") + article_dir.mkdir() + article_dir.joinpath("article.xml").write_bytes( + _make_article(f"

Results.

{_TABLE}") + ) + _articles._extract_tables(article_dir) + return article_dir def test_text_extractor_transform_failure(monkeypatch): @@ -12,3 +71,94 @@ def test_text_extractor_transform_failure(monkeypatch): etree, "XSLT", Mock(return_value=Mock(side_effect=ValueError)) ) assert extractor.extract(Mock(), Mock(), {}) == {} + + +def test_text_extractor_preserves_xref_text_by_default(): + extractor = _text.TextExtractor() + article = etree.fromstring( + _make_article( + "

Example ( Doe et al. ) test.

" + ) + ) + result = extractor.extract(article, Path("."), {}) + body = " ".join(result["body"].split()) + assert "Doe et al." in body + assert "( )" not in body + + +def test_text_extractor_strips_xref_text_when_disabled(): + extractor = _text.TextExtractor(preserve_cross_references=False) + article = etree.fromstring( + _make_article( + "

Example ( Doe et al. ) test.

" + ) + ) + result = extractor.extract(article, Path("."), {}) + body = " ".join(result["body"].split()) + assert "Doe et al." not in body + assert "( )" in body + + +def test_text_extractor_omits_tables_by_default(article_with_table): + extractor = _text.TextExtractor() + article = etree.parse(str(article_with_table.joinpath("article.xml"))) + result = extractor.extract(article, article_with_table, {}) + assert "Peak coordinates." in result["body"] + assert "-42" not in result["body"] + + +def test_text_extractor_keeps_tables_when_enabled(article_with_table): + """Header rows, merged cells and the footer are kept with the table.""" + extractor = _text.TextExtractor(keep_tables=True) + article = etree.parse(str(article_with_table.joinpath("article.xml"))) + result = extractor.extract(article, article_with_table, {}) + body = result["body"] + assert _TABLE_TEXT in body + # the table is inserted where it appears in the article + assert body.index("Peak coordinates.") < body.index(_TABLE_TEXT) + assert "pubget-table" not in body + + +def test_table_text_reproduces_table_csv(tmp_path): + """Cells are not passed through pandas' type inference a second time.""" + tables_dir = tmp_path.joinpath("tables") + tables_dir.mkdir() + tables_dir.joinpath("table_000.csv").write_text( + "Unnamed: 0,Subject\n,0012\n", "UTF-8" + ) + tables_dir.joinpath("table_000_info.json").write_text( + json.dumps( + { + "table_label": None, + "table_foot": None, + "n_header_rows": 1, + "table_data_file": "table_000.csv", + } + ), + "UTF-8", + ) + assert _text._load_tables(tmp_path) == {0: "\tSubject\n\t0012\n"} + + +def test_text_extractor_drops_placeholders_of_missing_tables( + article_with_table, +): + """Tables that `extract_articles` failed to parse leave no placeholder.""" + 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 "pubget-table" not in result["body"] + + +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 "pubget-table" not in result["body"] + assert "failed to read table" in caplog.text