diff --git a/src/political_event_tracking_research/official_event_import.py b/src/political_event_tracking_research/official_event_import.py index cc78698..e1706d6 100644 --- a/src/political_event_tracking_research/official_event_import.py +++ b/src/political_event_tracking_research/official_event_import.py @@ -141,6 +141,9 @@ def normalize_records(records: list[OfficialRecord]) -> list[dict[str, str]]: "confidence": confidence_for_source(record), "source_url": record.source_url, "notes": record.summary, + "entity_match_type": record.entity_match_type, + "match_evidence": record.match_evidence, + "relationship_type": record.relationship_type, } ) rows.sort(key=lambda row: (row["event_date"], row["symbol"], row["event_id"])) @@ -152,7 +155,19 @@ def import_official_events(input_path: str | Path, output_path: str | Path) -> l rows = normalize_records(records) write_csv_rows( output_path, - ["event_id", "event_date", "symbol", "event_type", "direction", "confidence", "source_url", "notes"], + [ + "event_id", + "event_date", + "symbol", + "event_type", + "direction", + "confidence", + "source_url", + "notes", + "entity_match_type", + "match_evidence", + "relationship_type", + ], rows, ) return rows diff --git a/tests/test_official_event_import.py b/tests/test_official_event_import.py index 8f7fc92..895bcca 100644 --- a/tests/test_official_event_import.py +++ b/tests/test_official_event_import.py @@ -38,6 +38,9 @@ def test_import_official_events_normalizes_to_event_schema(tmp_path: Path) -> No "confidence", "source_url", "notes", + "entity_match_type", + "match_evidence", + "relationship_type", } @@ -137,3 +140,20 @@ def test_community_lead_records_are_not_in_stable_source_set() -> None: with pytest.raises(ValueError, match="unsupported source_type"): normalize_records([record]) + + +def test_normalize_records_preserves_verified_entity_relationship_fields(tmp_path: Path) -> None: + input_path = tmp_path / "entity-fields.csv" + input_path.write_text( + "record_id,record_date,symbol,source_type,event_type,direction,source_url,summary," + "entity_match_type,match_evidence,relationship_type\n" + "entity-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish," + "https://www.sec.gov/example/entity-1,Entity record.,issuer,SEC filing names EVT1,issuer\n", + encoding="utf-8", + ) + + rows = normalize_records(load_official_records(input_path)) + + assert rows[0]["entity_match_type"] == "issuer" + assert rows[0]["match_evidence"] == "SEC filing names EVT1" + assert rows[0]["relationship_type"] == "issuer"