diff --git a/src/corppa/utils/dataset_refine.py b/src/corppa/utils/dataset_refine.py new file mode 100644 index 00000000..f65e24d6 --- /dev/null +++ b/src/corppa/utils/dataset_refine.py @@ -0,0 +1,190 @@ +""" +Utility script for PPA full-text dataset publication prep. + +Takes a manually cleaned version of PPA metadata in CSV format, and +propagates refined author and pub_place fields to CSV and JSON metadata +in the specified ppa_corpus folder. +""" + +import argparse +import json +import shutil +from pathlib import Path + +import polars as pl + + +def refine_metadata( + refined_metadata: Path, csv_metadata_path: Path, json_metadata_path: Path +) -> None: + refined_metadata_df = pl.read_csv(refined_metadata) + # get the last update time of records in refined metadata (all records have added/updated time) + refined_last_update = refined_metadata_df["updated"].max() + df = pl.read_csv(csv_metadata_path) + csv_columns = df.columns # store columns original order for output file + csv_total = df.height + + # join with refined metadata on work id so we can propagate work-specific + # resolutions (author names are unambiguous but pub places are not) + refined_df = df.join( + refined_metadata_df.select("work_id", "author", "pub_place"), + on="work_id", + suffix="_refined", + how="left", # preserve all original records + ) + # report how many instances of each field will be updated. A change is + # counted only when the refined column carries a value that differs from + # the original (the original is filled so comparisons with null are safe); + # an empty refined value means "leave the original alone". + author_changes_df = refined_df.filter( + pl.col.author_refined.is_not_null() + & pl.col.author_refined.ne(pl.col.author.fill_null("")) + ) + num_author_changes = author_changes_df.height + uniq_author_changes = ( + author_changes_df.select("author", "author_refined").unique().height + ) + print( + f"Updated author in {num_author_changes:,} works ({uniq_author_changes:,} unique replacements)" + ) + + pubplace_changes_df = refined_df.filter( + pl.col.pub_place_refined.is_not_null() + & pl.col.pub_place_refined.ne(pl.col.pub_place.fill_null("")) + ) + num_pubplace_changes = pubplace_changes_df.height + uniq_pubplace_changes = pubplace_changes_df.select( + "pub_place", "pub_place_refined" + ).n_unique() + print( + f"Updated pub_place in {num_pubplace_changes:,} works ({uniq_pubplace_changes:,} unique replacements)" + ) + + # then replace the original with the refined fields, but preserve the original + # value whenever the refined set is absent + refined_df = refined_df.with_columns( + author=pl.coalesce("author_refined", "author"), + pub_place=pl.coalesce("pub_place_refined", "pub_place"), + ).drop("author_refined", "pub_place_refined") + + # identify any records added after the manual cleanup, and update author/pub_place + # if there is an unambiguous mapping from the refined set + post_refine_records = df.filter(pl.col.added.gt(refined_last_update)) + if post_refine_records.height: + author_lookup = author_changes_df.select("author", "author_refined").unique() + # omit any ambiguous authors (a single original author resolving to more + # than one refined value cannot be applied automatically, and would make + # the join below emit duplicate rows per work_id) + author_lookup = author_lookup.filter(~author_lookup["author"].is_duplicated()) + + pub_place_lookup = pubplace_changes_df.select( + "pub_place", "pub_place_refined" + ).unique() + # omit any ambiguous placenames (i.e. Cambridge and Rochester could either be US or UK) + pub_place_lookup = pub_place_lookup.filter( + ~pub_place_lookup["pub_place"].is_duplicated() + ) + # join and then filter to any that have changes + post_refine_records = ( + post_refine_records.select("work_id", "author", "pub_place") + .join(author_lookup, on="author", how="left") + .join(pub_place_lookup, on="pub_place", how="left") + .filter( + # limit to records with at least one refined value + ~(pl.col.author_refined.is_null() & pl.col.pub_place_refined.is_null()) + ) + ) + + # update the refined data with these changes + refined_df = ( + refined_df.join( + post_refine_records.select( + "work_id", "author_refined", "pub_place_refined" + ), + on="work_id", + how="left", + ) + .with_columns( + # take the refined value if not null, otherwise the previous value + author=pl.when(pl.col.author_refined.is_not_null()) + .then(pl.col.author_refined) + .otherwise(pl.col.author), + pub_place=pl.when(pl.col.pub_place_refined.is_not_null()) + .then(pl.col.pub_place_refined) + .otherwise(pl.col.pub_place), + ) + .drop("author_refined", "pub_place_refined") + ) + + # replace refined fields with unrefined and save the file + assert refined_df.height == csv_total # no records lost + + # Path.copy not available in 3.12, so use shutil to make a backup + shutil.copy(csv_metadata_path, csv_metadata_path.with_suffix(".csv.bk")) + # then replace the original with the refined version + refined_df.select(csv_columns).write_csv(csv_metadata_path) + + # load json, update based on csv, write out + # NOTE: the CSV and JSON are known to be equivalent and in the same order + # (both are generated from the same data), so rows are paired in sequence + # the per-row assert is a check to catch any mismatches + with json_metadata_path.open() as jsonfile: + json_metadata = json.load(jsonfile) + refined_json_metadata = [] + for csv_row, json_row in zip(refined_df.iter_rows(named=True), json_metadata): + assert csv_row["work_id"] == json_row["work_id"] + if "author" in json_row: + json_row["author"] = csv_row["author"] + if "pub_place" in json_row: + json_row["pub_place"] = csv_row["pub_place"] + refined_json_metadata.append(json_row) + # no rows lost + assert len(json_metadata) == len(refined_json_metadata) == csv_total + + # make a backup and output the refined version + shutil.copy(json_metadata_path, json_metadata_path.with_suffix(".json.bk")) + with json_metadata_path.open("w") as jsonfile: + json.dump(refined_json_metadata, jsonfile, indent=2) + + print( + f"Updated metadata has been saved to {csv_metadata_path} and {json_metadata_path}" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Propagate PPA metadata standardization to updated metadata files" + ) + parser.add_argument( + "refined_metadata", + help="PPA metadata file with standardized author/pub_place fields to propagate", + type=Path, + ) + + parser.add_argument( + "corpus_dir", + help="PPA full-text corpus directory; expected to contain " + "work-level metadata files ppa_metadata.csv and ppa_metadata.json", + type=Path, + ) + + args = parser.parse_args() + if not args.refined_metadata.is_file(): + raise SystemExit(f"Input metadata file {args.refined_metadata} does not exist") + + if not args.corpus_dir.is_dir(): + raise SystemExit(f"PPA corpus path {args.corpus_dir} is not a directory") + + # we expect two formats + csv_metadata_path = args.corpus_dir / "ppa_metadata.csv" + json_metadata_path = args.corpus_dir / "ppa_metadata.json" + if not csv_metadata_path.exists(): + raise SystemExit(f"PPA metadata file {csv_metadata_path} not found") + if not json_metadata_path.exists(): + raise SystemExit(f"PPA metadata file {json_metadata_path} not found") + + refine_metadata(args.refined_metadata, csv_metadata_path, json_metadata_path) + + +if __name__ == "__main__": + main() diff --git a/tests/test_utils/test_dataset_refine.py b/tests/test_utils/test_dataset_refine.py new file mode 100644 index 00000000..d3eaf9df --- /dev/null +++ b/tests/test_utils/test_dataset_refine.py @@ -0,0 +1,401 @@ +# Copyright (c) 2024-2026, Center for Digital Humanities, Princeton University +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path + +import polars as pl +import pytest + +import corppa.utils.dataset_refine as dataset_refine + + +def _make_corpus(tmp_path: Path, rows, refined_rows) -> tuple[Path, Path]: + """Write ppa_metadata.csv/json + refined.csv; returns (corpus_dir, refined_csv).""" + corpus_dir = tmp_path / "corpus" + corpus_dir.mkdir() + pl.DataFrame(rows).write_csv(corpus_dir / "ppa_metadata.csv") + refined_path = tmp_path / "refined.csv" + pl.DataFrame(refined_rows).write_csv(refined_path) + with (corpus_dir / "ppa_metadata.json").open("w") as f: + json.dump(rows, f, indent=2) + return corpus_dir, refined_path + + +def _csv_columns(corpus_dir: Path) -> dict[str, list]: + return pl.read_csv(corpus_dir / "ppa_metadata.csv").to_dict(as_series=False) + + +def _json_rows(corpus_dir: Path) -> list[dict]: + with (corpus_dir / "ppa_metadata.json").open() as f: + return json.load(f) + + +def test_propagates_refined_author_and_pub_place(tmp_path): + rows = [ + { + "work_id": "w1", + "author": "Alice A.", + "pub_place": "Cambridge", + "added": "2024-01-01", + }, + { + "work_id": "w2", + "author": "Bob B.", + "pub_place": "Oxford", + "added": "2024-01-01", + }, + ] + refined_rows = [ + { + "work_id": "w1", + "author": "Alice Adams", + "pub_place": "Cambridge (UK)", + "updated": "2024-02-01", + }, + { + "work_id": "w2", + "author": "Bob Brown", + "pub_place": "Oxford", + "updated": "2024-02-01", + }, + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + out = _csv_columns(corpus_dir) + assert out["author"] == ["Alice Adams", "Bob Brown"] + assert out["pub_place"] == ["Cambridge (UK)", "Oxford"] + js = _json_rows(corpus_dir) + assert js[0]["author"] == "Alice Adams" + assert js[0]["pub_place"] == "Cambridge (UK)" + assert "added" in out # extra columns are preserved + + +def test_unmatched_record_keeps_original_values(tmp_path): + # a work absent from the refined set must not have its author/pub_place + # wiped to null + rows = [ + { + "work_id": "w1", + "author": "Alice A.", + "pub_place": "Cambridge", + "added": "2024-01-01", + }, + { + "work_id": "w2", + "author": "Bob B.", + "pub_place": "Oxford", + "added": "2024-01-01", + }, + ] + refined_rows = [ + { + "work_id": "w1", + "author": "Alice Adams", + "pub_place": "Cambridge (UK)", + "updated": "2024-02-01", + }, + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + out = _csv_columns(corpus_dir) + assert out["author"] == ["Alice Adams", "Bob B."] + assert out["pub_place"] == ["Cambridge (UK)", "Oxford"] + js = _json_rows(corpus_dir) + assert js[1]["author"] == "Bob B." + + +def test_fills_null_original_author_from_refined(tmp_path): + rows = [ + {"work_id": "w1", "author": None, "pub_place": "Boston", "added": "2024-01-01"} + ] + refined_rows = [ + { + "work_id": "w1", + "author": "Anonymous", + "pub_place": "Boston", + "updated": "2024-02-01", + } + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + assert _csv_columns(corpus_dir)["author"] == ["Anonymous"] + + +def test_post_refine_records_updated_via_unambiguous_lookup(tmp_path): + rows = [ + { + "work_id": "w1", + "author": "Doe", + "pub_place": "Cambridge", + "added": "2024-01-01", + }, + { + "work_id": "w2", + "author": "Doe", + "pub_place": "Cambridge", + "added": "2024-07-01", + }, + ] + refined_rows = [ + { + "work_id": "w1", + "author": "John Doe", + "pub_place": "Cambridge (UK)", + "updated": "2024-06-01", + }, + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + out = _csv_columns(corpus_dir) + assert out["author"] == ["John Doe", "John Doe"] + assert out["pub_place"] == ["Cambridge (UK)", "Cambridge (UK)"] + + +def test_ambiguous_author_lookup_is_not_applied(tmp_path): + # "Doe" resolves to two different refined authors, so a post-refine record + # with author "Doe" cannot be updated automatically and keeps its value + rows = [ + {"work_id": "w1", "author": "Doe", "pub_place": "X", "added": "2024-01-01"}, + {"work_id": "w2", "author": "Doe", "pub_place": "Y", "added": "2024-01-01"}, + {"work_id": "w3", "author": "Doe", "pub_place": "Z", "added": "2024-07-01"}, + ] + refined_rows = [ + { + "work_id": "w1", + "author": "John Doe", + "pub_place": "X", + "updated": "2024-06-01", + }, + { + "work_id": "w2", + "author": "Jane Doe", + "pub_place": "Y", + "updated": "2024-06-01", + }, + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + out = _csv_columns(corpus_dir) + assert out["author"] == ["John Doe", "Jane Doe", "Doe"] + # no rows were lost or duplicated by the ambiguous lookup + assert len(out["author"]) == 3 + + +def test_ambiguous_pub_place_lookup_is_not_applied(tmp_path): + # "Cambridge" resolves to two distinct refined places; a post-refine record + # with pub_place "Cambridge" keeps its original value + rows = [ + { + "work_id": "w1", + "author": "A", + "pub_place": "Cambridge", + "added": "2024-01-01", + }, + { + "work_id": "w2", + "author": "B", + "pub_place": "Cambridge", + "added": "2024-01-01", + }, + { + "work_id": "w3", + "author": "C", + "pub_place": "Cambridge", + "added": "2024-07-01", + }, + ] + refined_rows = [ + { + "work_id": "w1", + "author": "A", + "pub_place": "Cambridge (US)", + "updated": "2024-06-01", + }, + { + "work_id": "w2", + "author": "B", + "pub_place": "Cambridge (UK)", + "updated": "2024-06-01", + }, + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + out = _csv_columns(corpus_dir) + assert out["pub_place"] == ["Cambridge (US)", "Cambridge (UK)", "Cambridge"] + assert len(out["pub_place"]) == 3 + + +def test_duplicate_work_id_in_refined_raises(tmp_path): + # a duplicated work_id in the refined set inflates the join and trips the + # no-records-lost assertion + rows = [{"work_id": "w1", "author": "A", "pub_place": "X", "added": "2024-01-01"}] + refined_rows = [ + {"work_id": "w1", "author": "B", "pub_place": "X", "updated": "2024-02-01"}, + {"work_id": "w1", "author": "C", "pub_place": "X", "updated": "2024-02-01"}, + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + with pytest.raises(AssertionError): + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + +def test_backup_files_created(tmp_path): + rows = [{"work_id": "w1", "author": "A", "pub_place": "X", "added": "2024-01-01"}] + refined_rows = [ + {"work_id": "w1", "author": "B", "pub_place": "Y", "updated": "2024-02-01"} + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + assert (corpus_dir / "ppa_metadata.csv.bk").is_file() + assert (corpus_dir / "ppa_metadata.json.bk").is_file() + # backups carry the original, unrefined values + assert pl.read_csv(corpus_dir / "ppa_metadata.csv.bk")["author"].to_list() == ["A"] + + +def test_column_order_and_row_count_preserved(tmp_path): + rows = [ + { + "work_id": "w1", + "title": "T1", + "author": "A", + "pub_place": "X", + "added": "2024-01-01", + }, + { + "work_id": "w2", + "title": "T2", + "author": "B", + "pub_place": "Y", + "added": "2024-01-01", + }, + ] + refined_rows = [ + {"work_id": "w1", "author": "C", "pub_place": "Z", "updated": "2024-02-01"} + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + csv_path = corpus_dir / "ppa_metadata.csv" + original_columns = pl.read_csv(csv_path).columns + + dataset_refine.refine_metadata(refined, csv_path, corpus_dir / "ppa_metadata.json") + + out_df = pl.read_csv(csv_path) + assert out_df.columns == original_columns + assert out_df.height == 2 + + +def test_json_row_without_author_pub_place_keys_unchanged(tmp_path): + rows = [ + {"work_id": "w1", "author": "A", "pub_place": "X", "added": "2024-01-01"}, + {"work_id": "w2", "title": "no author field", "added": "2024-01-01"}, + ] + refined_rows = [ + {"work_id": "w1", "author": "B", "pub_place": "Y", "updated": "2024-02-01"} + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + js = _json_rows(corpus_dir) + assert js[0]["author"] == "B" + # the second json row has neither key, so it stays untouched even though its + # csv counterpart's author was (coalesced)... + assert "author" not in js[1] + assert js[1]["title"] == "no author field" + + +def test_json_order_mismatch_raises(tmp_path): + rows = [ + {"work_id": "w1", "author": "A", "pub_place": "X", "added": "2024-01-01"}, + {"work_id": "w2", "author": "B", "pub_place": "Y", "added": "2024-01-01"}, + ] + refined_rows = [ + {"work_id": "w1", "author": "C", "pub_place": "Z", "updated": "2024-02-01"} + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + # shuffle the json rows so they no longer match the csv order + with (corpus_dir / "ppa_metadata.json").open("w") as f: + json.dump([rows[1], rows[0]], f, indent=2) + + with pytest.raises(AssertionError): + dataset_refine.refine_metadata( + refined, corpus_dir / "ppa_metadata.csv", corpus_dir / "ppa_metadata.json" + ) + + +# --- main() CLI --- + + +def test_main_missing_refined_file_exits(tmp_path, monkeypatch): + corpus_dir = tmp_path / "corpus" + corpus_dir.mkdir() + monkeypatch.setattr( + "sys.argv", + ["dataset_refine.py", str(tmp_path / "nope.csv"), str(corpus_dir)], + ) + + with pytest.raises(SystemExit, match="Input metadata file"): + dataset_refine.main() + + +def test_main_missing_csv_exits(tmp_path, monkeypatch): + corpus_dir = tmp_path / "corpus" + corpus_dir.mkdir() + refined = tmp_path / "refined.csv" + refined.write_text("work_id,author,pub_place,updated\n") + with (corpus_dir / "ppa_metadata.json").open("w") as f: + json.dump([], f) # only json present; csv is missing + monkeypatch.setattr( + "sys.argv", ["dataset_refine.py", str(refined), str(corpus_dir)] + ) + + with pytest.raises(SystemExit, match="ppa_metadata.csv not found"): + dataset_refine.main() + + +def test_main_full_run_exits_zero(tmp_path, monkeypatch, capsys): + rows = [{"work_id": "w1", "author": "A", "pub_place": "X", "added": "2024-01-01"}] + refined_rows = [ + {"work_id": "w1", "author": "B", "pub_place": "Y", "updated": "2024-02-01"} + ] + corpus_dir, refined = _make_corpus(tmp_path, rows, refined_rows) + monkeypatch.setattr( + "sys.argv", ["dataset_refine.py", str(refined), str(corpus_dir)] + ) + + dataset_refine.main() + + assert "Updated author in 1 works" in capsys.readouterr().out + assert _csv_columns(corpus_dir)["author"] == ["B"]