Support current region scoping metadata (#364)

Author: MaxGhenis

changelog.d/364.fixed.md

@@ -0,0 +1 @@
+Support current PolicyEngine region scoping metadata when reseeding regions.

scripts/seed_regions.py

@@ -21,6 +21,7 @@
 
 import argparse
 import time
+from typing import Any
 
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from seed_utils import console, get_session
@@ -36,6 +37,34 @@
 from policyengine_api.models.region import RegionType  # noqa: E402
 
 
+def _region_scoping_strategy(pe_region: Any) -> Any | None:
+    return getattr(pe_region, "scoping_strategy", None)
+
+
+def _region_requires_filter(pe_region: Any) -> bool:
+    strategy = _region_scoping_strategy(pe_region)
+    return bool(getattr(pe_region, "requires_filter", strategy is not None))
+
+
+def _region_filter_field(pe_region: Any) -> str | None:
+    if hasattr(pe_region, "filter_field"):
+        return pe_region.filter_field
+    strategy = _region_scoping_strategy(pe_region)
+    return getattr(strategy, "variable_name", None)
+
+
+def _region_filter_value(pe_region: Any) -> str | None:
+    if hasattr(pe_region, "filter_value"):
+        return pe_region.filter_value
+    strategy = _region_scoping_strategy(pe_region)
+    return getattr(strategy, "variable_value", None)
+
+
+def _region_filter_strategy(pe_region: Any) -> str | None:
+    strategy = _region_scoping_strategy(pe_region)
+    return getattr(strategy, "strategy_type", None)
+
+
 def _group_us_datasets(
     session: Session,
     us_model_id,
@@ -197,14 +226,10 @@ def seed_us_regions(
                     code=pe_region.code,
                     label=pe_region.label,
                     region_type=RegionType(pe_region.region_type),
-                    requires_filter=pe_region.requires_filter,
-                    filter_field=pe_region.filter_field,
-                    filter_value=pe_region.filter_value,
-                    filter_strategy=(
-                        pe_region.scoping_strategy.strategy_type
-                        if pe_region.scoping_strategy
-                        else None
-                    ),
+                    requires_filter=_region_requires_filter(pe_region),
+                    filter_field=_region_filter_field(pe_region),
+                    filter_value=_region_filter_value(pe_region),
+                    filter_strategy=_region_filter_strategy(pe_region),
                     parent_code=pe_region.parent_code,
                     state_code=pe_region.state_code,
                     state_name=pe_region.state_name,
@@ -295,14 +320,10 @@ def seed_uk_regions(session: Session) -> tuple[int, int, int]:
                     code=pe_region.code,
                     label=pe_region.label,
                     region_type=RegionType(pe_region.region_type),
-                    requires_filter=pe_region.requires_filter,
-                    filter_field=pe_region.filter_field,
-                    filter_value=pe_region.filter_value,
-                    filter_strategy=(
-                        pe_region.scoping_strategy.strategy_type
-                        if pe_region.scoping_strategy
-                        else None
-                    ),
+                    requires_filter=_region_requires_filter(pe_region),
+                    filter_field=_region_filter_field(pe_region),
+                    filter_value=_region_filter_value(pe_region),
+                    filter_strategy=_region_filter_strategy(pe_region),
                     parent_code=pe_region.parent_code,
                     state_code=None,
                     state_name=None,

tests/test_seed_regions.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts"
+sys.path.insert(0, str(SCRIPTS_DIR))
+
+import seed_regions  # noqa: E402
+
+
+def test_region_filter_helpers_support_scoping_strategy_schema() -> None:
+    region = SimpleNamespace(
+        scoping_strategy=SimpleNamespace(
+            strategy_type="row_filter",
+            variable_name="place_fips",
+            variable_value="03000",
+        )
+    )
+
+    assert seed_regions._region_requires_filter(region) is True
+    assert seed_regions._region_filter_field(region) == "place_fips"
+    assert seed_regions._region_filter_value(region) == "03000"
+    assert seed_regions._region_filter_strategy(region) == "row_filter"
+
+
+def test_region_filter_helpers_support_legacy_filter_fields() -> None:
+    region = SimpleNamespace(
+        requires_filter=True,
+        filter_field="country",
+        filter_value="ENGLAND",
+        scoping_strategy=None,
+    )
+
+    assert seed_regions._region_requires_filter(region) is True
+    assert seed_regions._region_filter_field(region) == "country"
+    assert seed_regions._region_filter_value(region) == "ENGLAND"
+    assert seed_regions._region_filter_strategy(region) is None
+
+
+def test_current_policyengine_regions_do_not_require_legacy_filter_attrs() -> None:
+    from policyengine.countries.us.regions import us_region_registry
+
+    place = next(
+        region for region in us_region_registry.regions if region.region_type == "place"
+    )
+
+    assert not hasattr(place, "filter_field")
+    assert seed_regions._region_requires_filter(place) is True
+    assert seed_regions._region_filter_field(place) == "place_fips"
+    assert seed_regions._region_filter_value(place)
+    assert seed_regions._region_filter_strategy(place) == "row_filter"