Skip to content

Commit 45642af

Browse files
Pigbibicodex
andcommitted
feat: add official record entity matching fields
Co-Authored-By: Codex <noreply@openai.com>
1 parent 928d4f5 commit 45642af

2 files changed

Lines changed: 77 additions & 0 deletions

File tree

src/political_event_tracking_research/official_event_import.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,10 @@
3434
}
3535
)
3636

37+
ALLOWED_ENTITY_RELATIONSHIP_TYPES = frozenset(
38+
{"issuer", "direct_beneficiary", "industry_context", "unverified"}
39+
)
40+
3741

3842
@dataclass(frozen=True)
3943
class OfficialRecord:
@@ -45,6 +49,9 @@ class OfficialRecord:
4549
direction: str
4650
source_url: str
4751
summary: str
52+
entity_match_type: str = "unverified"
53+
match_evidence: str = ""
54+
relationship_type: str = "unverified"
4855

4956

5057
def slug(value: str) -> str:
@@ -76,6 +83,9 @@ def load_official_records(path: str | Path) -> list[OfficialRecord]:
7683
direction=row.get("direction", ""),
7784
source_url=row["source_url"],
7885
summary=row.get("summary", ""),
86+
entity_match_type=row.get("entity_match_type", "unverified"),
87+
match_evidence=row.get("match_evidence", ""),
88+
relationship_type=row.get("relationship_type", "unverified"),
7989
)
8090
)
8191
return records
@@ -87,6 +97,12 @@ def validate_record(record: OfficialRecord) -> None:
8797
raise ValueError("record_id is required")
8898
if not record.symbol.strip():
8999
raise ValueError(f"{record.record_id}: symbol is required")
100+
for field_name, value in (
101+
("entity_match_type", record.entity_match_type),
102+
("relationship_type", record.relationship_type),
103+
):
104+
if value not in ALLOWED_ENTITY_RELATIONSHIP_TYPES:
105+
raise ValueError(f"{record.record_id}: unsupported {field_name} {value!r}")
90106
if record.event_type not in ALLOWED_EVENT_TYPES:
91107
raise ValueError(f"{record.record_id}: unsupported event_type {record.event_type!r}")
92108
if record.source_type in GOVERNMENT_SOURCE_TYPES:

tests/test_official_event_import.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from political_event_tracking_research.official_event_import import (
88
OfficialRecord,
99
import_official_events,
10+
load_official_records,
1011
normalize_records,
1112
)
1213

@@ -28,6 +29,66 @@ def test_import_official_events_normalizes_to_event_schema(tmp_path: Path) -> No
2829
assert rows[-2]["event_id"] == "official-issuer-release-demo-issuer-evt3"
2930
assert rows[-1]["event_id"] == "official-financial-media-demo-media-evt5"
3031
assert rows[-1]["confidence"] == "low"
32+
assert set(rows[0]) == {
33+
"event_id",
34+
"event_date",
35+
"symbol",
36+
"event_type",
37+
"direction",
38+
"confidence",
39+
"source_url",
40+
"notes",
41+
}
42+
43+
44+
def test_load_official_records_defaults_entity_fields_for_legacy_csv(tmp_path: Path) -> None:
45+
input_path = tmp_path / "legacy.csv"
46+
input_path.write_text(
47+
"record_id,record_date,symbol,source_type,event_type,direction,source_url,summary\n"
48+
"legacy-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish,"
49+
"https://www.sec.gov/example/legacy-1,Legacy record.\n",
50+
encoding="utf-8",
51+
)
52+
53+
record = load_official_records(input_path)[0]
54+
55+
assert record.entity_match_type == "unverified"
56+
assert record.match_evidence == ""
57+
assert record.relationship_type == "unverified"
58+
59+
60+
@pytest.mark.parametrize("value", ["issuer", "direct_beneficiary", "industry_context", "unverified"])
61+
def test_new_entity_fields_accept_allowed_values(tmp_path: Path, value: str) -> None:
62+
input_path = tmp_path / "entity-fields.csv"
63+
input_path.write_text(
64+
"record_id,record_date,symbol,source_type,event_type,direction,source_url,summary,"
65+
"entity_match_type,match_evidence,relationship_type\n"
66+
"entity-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish,"
67+
f"https://www.sec.gov/example/entity-1,Entity record.,{value},SEC filing,{value}\n",
68+
encoding="utf-8",
69+
)
70+
71+
rows = normalize_records(load_official_records(input_path))
72+
73+
assert rows[0]["event_id"] == "official-government-filing-entity-1"
74+
75+
76+
@pytest.mark.parametrize("field", ["entity_match_type", "relationship_type"])
77+
@pytest.mark.parametrize("value", ["", " ", "not-allowed"])
78+
def test_new_entity_fields_fail_closed_for_blank_or_invalid_values(tmp_path: Path, field: str, value: str) -> None:
79+
input_path = tmp_path / f"invalid-{field}.csv"
80+
entity_match_type = value if field == "entity_match_type" else "issuer"
81+
relationship_type = value if field == "relationship_type" else "issuer"
82+
input_path.write_text(
83+
"record_id,record_date,symbol,source_type,event_type,direction,source_url,summary,"
84+
"entity_match_type,match_evidence,relationship_type\n"
85+
f"entity-1,2026-01-10,EVT1,government_filing,disclosure_buy,bullish,"
86+
f"https://www.sec.gov/example/entity-1,Entity record.,{entity_match_type},SEC filing,{relationship_type}\n",
87+
encoding="utf-8",
88+
)
89+
90+
with pytest.raises(ValueError):
91+
normalize_records(load_official_records(input_path))
3192

3293

3394
def test_government_records_reject_non_gov_urls() -> None:

0 commit comments

Comments
 (0)