From 45642af4bb2b6518e0dedcb80281540d4dca9670 Mon Sep 17 00:00:00 2001 From: Pigbibi <20649888+Pigbibi@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:36:49 +0800 Subject: [PATCH] feat: add official record entity matching fields Co-Authored-By: Codex --- .../official_event_import.py | 16 +++++ tests/test_official_event_import.py | 61 +++++++++++++++++++ 2 files changed, 77 insertions(+) diff --git a/src/political_event_tracking_research/official_event_import.py b/src/political_event_tracking_research/official_event_import.py index 947be00..cc78698 100644 --- a/src/political_event_tracking_research/official_event_import.py +++ b/src/political_event_tracking_research/official_event_import.py @@ -34,6 +34,10 @@ } ) +ALLOWED_ENTITY_RELATIONSHIP_TYPES = frozenset( + {"issuer", "direct_beneficiary", "industry_context", "unverified"} +) + @dataclass(frozen=True) class OfficialRecord: @@ -45,6 +49,9 @@ class OfficialRecord: direction: str source_url: str summary: str + entity_match_type: str = "unverified" + match_evidence: str = "" + relationship_type: str = "unverified" def slug(value: str) -> str: @@ -76,6 +83,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", "unverified"), + match_evidence=row.get("match_evidence", ""), + relationship_type=row.get("relationship_type", "unverified"), ) ) return records @@ -87,6 +97,12 @@ def validate_record(record: OfficialRecord) -> None: raise ValueError("record_id is required") if not record.symbol.strip(): raise ValueError(f"{record.record_id}: symbol is required") + for field_name, value in ( + ("entity_match_type", record.entity_match_type), + ("relationship_type", record.relationship_type), + ): + if value not in ALLOWED_ENTITY_RELATIONSHIP_TYPES: + raise ValueError(f"{record.record_id}: unsupported {field_name} {value!r}") if record.event_type not in ALLOWED_EVENT_TYPES: raise ValueError(f"{record.record_id}: unsupported event_type {record.event_type!r}") if record.source_type in GOVERNMENT_SOURCE_TYPES: diff --git a/tests/test_official_event_import.py b/tests/test_official_event_import.py index 5c315ef..8f7fc92 100644 --- a/tests/test_official_event_import.py +++ b/tests/test_official_event_import.py @@ -7,6 +7,7 @@ from political_event_tracking_research.official_event_import import ( OfficialRecord, import_official_events, + load_official_records, normalize_records, ) @@ -28,6 +29,66 @@ def test_import_official_events_normalizes_to_event_schema(tmp_path: Path) -> No assert rows[-2]["event_id"] == "official-issuer-release-demo-issuer-evt3" assert rows[-1]["event_id"] == "official-financial-media-demo-media-evt5" assert rows[-1]["confidence"] == "low" + assert set(rows[0]) == { + "event_id", + "event_date", + "symbol", + "event_type", + "direction", + "confidence", + "source_url", + "notes", + } + + +def test_load_official_records_defaults_entity_fields_for_legacy_csv(tmp_path: Path) -> None: + input_path = tmp_path / "legacy.csv" + input_path.write_text( + "record_id,record_date,symbol,source_type,event_type,direction,source_url,summary\n" + "legacy-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish," + "https://www.sec.gov/example/legacy-1,Legacy record.\n", + encoding="utf-8", + ) + + record = load_official_records(input_path)[0] + + assert record.entity_match_type == "unverified" + assert record.match_evidence == "" + assert record.relationship_type == "unverified" + + +@pytest.mark.parametrize("value", ["issuer", "direct_beneficiary", "industry_context", "unverified"]) +def test_new_entity_fields_accept_allowed_values(tmp_path: Path, value: str) -> 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," + f"https://www.sec.gov/example/entity-1,Entity record.,{value},SEC filing,{value}\n", + encoding="utf-8", + ) + + rows = normalize_records(load_official_records(input_path)) + + assert rows[0]["event_id"] == "official-government-filing-entity-1" + + +@pytest.mark.parametrize("field", ["entity_match_type", "relationship_type"]) +@pytest.mark.parametrize("value", ["", " ", "not-allowed"]) +def test_new_entity_fields_fail_closed_for_blank_or_invalid_values(tmp_path: Path, field: str, value: str) -> None: + input_path = tmp_path / f"invalid-{field}.csv" + entity_match_type = value if field == "entity_match_type" else "issuer" + relationship_type = value if field == "relationship_type" else "issuer" + 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" + f"entity-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish," + f"https://www.sec.gov/example/entity-1,Entity record.,{entity_match_type},SEC filing,{relationship_type}\n", + encoding="utf-8", + ) + + with pytest.raises(ValueError): + normalize_records(load_official_records(input_path)) def test_government_records_reject_non_gov_urls() -> None: