Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 46 additions & 3 deletions src/political_event_tracking_research/official_event_import.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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")
Expand Down Expand Up @@ -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
Expand All @@ -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__":
Expand Down
120 changes: 120 additions & 0 deletions tests/test_official_event_import.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import csv
from pathlib import Path

import pytest
Expand Down Expand Up @@ -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",
Expand Down
Loading