diff --git a/src/political_event_tracking_research/official_event_import.py b/src/political_event_tracking_research/official_event_import.py index 947be00..8e09cea 100644 --- a/src/political_event_tracking_research/official_event_import.py +++ b/src/political_event_tracking_research/official_event_import.py @@ -34,6 +34,26 @@ } ) +ALLOWED_ENTITY_MATCH_TYPES = frozenset( + {"issuer", "direct_beneficiary", "industry_context", "unverified"} +) + +LEGACY_EVENT_COLUMNS = [ + "event_id", + "event_date", + "symbol", + "event_type", + "direction", + "confidence", + "source_url", + "notes", +] +ENTITY_EVENT_COLUMNS = LEGACY_EVENT_COLUMNS + [ + "entity_match_type", + "match_evidence", + "relationship_type", +] + @dataclass(frozen=True) class OfficialRecord: @@ -45,6 +65,9 @@ class OfficialRecord: direction: str source_url: str summary: str + entity_match_type: str = "unverified" + match_evidence: str = "" + relationship_type: str = "" def slug(value: str) -> str: @@ -76,6 +99,9 @@ def load_official_records(path: str | Path) -> list[OfficialRecord]: direction=row.get("direction", ""), source_url=row["source_url"], summary=row.get("summary", ""), + entity_match_type=row.get("entity_match_type") or "unverified", + match_evidence=row.get("match_evidence", ""), + relationship_type=row.get("relationship_type", ""), ) ) return records @@ -89,6 +115,10 @@ def validate_record(record: OfficialRecord) -> None: raise ValueError(f"{record.record_id}: symbol is required") if record.event_type not in ALLOWED_EVENT_TYPES: raise ValueError(f"{record.record_id}: unsupported event_type {record.event_type!r}") + if record.entity_match_type not in ALLOWED_ENTITY_MATCH_TYPES: + raise ValueError( + f"{record.record_id}: unsupported entity_match_type {record.entity_match_type!r}" + ) if record.source_type in GOVERNMENT_SOURCE_TYPES: if not is_government_url(record.source_url): raise ValueError(f"{record.record_id}: government source URLs must be https .gov URLs") @@ -125,18 +155,26 @@ 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"])) return rows -def import_official_events(input_path: str | Path, output_path: str | Path) -> list[dict[str, str]]: +def import_official_events( + input_path: str | Path, + output_path: str | Path, + *, + include_entity_metadata: bool = False, +) -> list[dict[str, str]]: records = load_official_records(input_path) rows = normalize_records(records) write_csv_rows( output_path, - ["event_id", "event_date", "symbol", "event_type", "direction", "confidence", "source_url", "notes"], + ENTITY_EVENT_COLUMNS if include_entity_metadata else LEGACY_EVENT_COLUMNS, rows, ) return rows @@ -148,12 +186,17 @@ def build_arg_parser() -> argparse.ArgumentParser: ) parser.add_argument("--input", required=True, help="Official-source record CSV.") parser.add_argument("--output", required=True, help="Output normalized event CSV.") + parser.add_argument( + "--include-entity-metadata", + action="store_true", + help="Opt in to entity metadata columns in the normalized CSV.", + ) return parser def main(argv: list[str] | None = None) -> None: args = build_arg_parser().parse_args(argv) - import_official_events(args.input, args.output) + import_official_events(args.input, args.output, include_entity_metadata=args.include_entity_metadata) if __name__ == "__main__": diff --git a/tests/test_official_event_import.py b/tests/test_official_event_import.py index 5c315ef..476258c 100644 --- a/tests/test_official_event_import.py +++ b/tests/test_official_event_import.py @@ -1,5 +1,6 @@ from __future__ import annotations +import csv from pathlib import Path import pytest @@ -30,6 +31,125 @@ def test_import_official_events_normalizes_to_event_schema(tmp_path: Path) -> No assert rows[-1]["confidence"] == "low" +def test_legacy_input_defaults_to_legacy_eight_column_serialization(tmp_path: Path) -> None: + output = tmp_path / "events.csv" + + rows = import_official_events(ROOT / "examples/official_records.example.csv", output) + + assert all(row["entity_match_type"] == "unverified" for row in rows) + assert all(row["match_evidence"] == "" for row in rows) + assert all(row["relationship_type"] == "" for row in rows) + with output.open(newline="", encoding="utf-8") as handle: + assert next(csv.reader(handle)) == [ + "event_id", + "event_date", + "symbol", + "event_type", + "direction", + "confidence", + "source_url", + "notes", + ] + + +def test_entity_metadata_serialization_requires_explicit_opt_in(tmp_path: Path) -> None: + input_path = tmp_path / "official-records.csv" + output = tmp_path / "events.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-record,2026-01-10,EVT,government_filing,disclosure_buy,bullish," + "https://www.sec.gov/example/entity-record,Entity evidence.,direct_beneficiary," + "Named beneficiary in filing.,supplier\n", + encoding="utf-8", + ) + + import_official_events(input_path, output, include_entity_metadata=True) + + with output.open(newline="", encoding="utf-8") as handle: + rows = list(csv.DictReader(handle)) + assert list(rows[0]) == [ + "event_id", + "event_date", + "symbol", + "event_type", + "direction", + "confidence", + "source_url", + "notes", + "entity_match_type", + "match_evidence", + "relationship_type", + ] + assert rows[0]["entity_match_type"] == "direct_beneficiary" + assert rows[0]["match_evidence"] == "Named beneficiary in filing." + assert rows[0]["relationship_type"] == "supplier" + + +def test_entity_schema_fields_are_imported_and_normalized() -> None: + record = OfficialRecord( + record_id="entity-record", + record_date="2026-01-10", + symbol="EVT", + source_type="government_filing", + event_type="disclosure_buy", + direction="bullish", + source_url="https://www.sec.gov/example/entity-record", + summary="Entity evidence.", + entity_match_type="direct_beneficiary", + match_evidence="Named beneficiary in filing.", + relationship_type="supplier", + ) + + assert normalize_records([record])[0] == { + "event_id": "official-government-filing-entity-record", + "event_date": "2026-01-10", + "symbol": "EVT", + "event_type": "disclosure_buy", + "direction": "bullish", + "confidence": "high", + "source_url": "https://www.sec.gov/example/entity-record", + "notes": "Entity evidence.", + "entity_match_type": "direct_beneficiary", + "match_evidence": "Named beneficiary in filing.", + "relationship_type": "supplier", + } + + +@pytest.mark.parametrize("entity_match_type", ["issuer", "direct_beneficiary", "industry_context", "unverified"]) +def test_entity_match_type_accepts_only_contract_values(entity_match_type: str) -> None: + record = OfficialRecord( + record_id="entity-record", + record_date="2026-01-10", + symbol="EVT", + source_type="government_filing", + event_type="disclosure_buy", + direction="bullish", + source_url="https://www.sec.gov/example/entity-record", + summary="Entity evidence.", + entity_match_type=entity_match_type, + ) + + normalize_records([record]) + + +def test_entity_match_type_rejects_unknown_values() -> None: + record = OfficialRecord( + record_id="entity-record", + record_date="2026-01-10", + symbol="EVT", + source_type="government_filing", + event_type="disclosure_buy", + direction="bullish", + source_url="https://www.sec.gov/example/entity-record", + summary="Entity evidence.", + entity_match_type="inferred", + ) + + with pytest.raises(ValueError, match="unsupported entity_match_type"): + normalize_records([record]) + + def test_government_records_reject_non_gov_urls() -> None: record = OfficialRecord( record_id="bad-media-record",