fix: Clean up function, changelog

Author: anth-volk

changelog_entry.yaml

@@ -1 +1,8 @@
-
+- bump: patch
+  changes:
+    added:
+    - Conversion of county FIPS codes to county enum items
+    - Helper function to convert string county names to enum keys
+    changed:
+    - Modified county variable to depend on FIPS input, then on ZIP code
+    - Modified county variable to use helper function for conversion from county names to enum keys

policyengine_us/tools/geography/county_helpers.py

@@ -0,0 +1,25 @@
+import pandas as pd
+import numpy as np
+from policyengine_us.variables.household.demographic.geographic.county.county_enum import (
+    County,
+)
+
+
+def map_county_string_to_enum(
+    county_name: "pd.Series[str]", state_code: "pd.Series[str]"
+) -> "pd.Series[int]":
+    """Helper function to map county name and state code to County enum value."""
+    county_key = county_name.apply(
+        lambda name: name.replace(" ", "_")
+        .replace("-", "_")
+        .replace(".", "")
+        .replace("'", "_")
+        .strip()
+        .upper()
+    )
+    county_state = county_key.str.cat(state_code, sep="_")
+    county_names = pd.Series(
+        np.arange(len(County._member_names_)),
+        index=County._member_names_,
+    )
+    return county_names[county_state]

policyengine_us/variables/household/demographic/geographic/county/county.py

@@ -1,4 +1,7 @@
 from policyengine_us.model_api import *
+from policyengine_us.tools.geography.county_helpers import (
+    map_county_string_to_enum,
+)
 from policyengine_us.variables.household.demographic.geographic.county.county_enum import (
     County,
 )
@@ -21,36 +24,11 @@ def formula(household, period, parameters):
             county_fips_codes = COUNTY_FIPS_DATASET.set_index("county_fips")
             county_name = county_fips_codes.loc[county_fips, "county_name"]
             state_code = county_fips_codes.loc[county_fips, "state"]
-            county_key = county_name.apply(
-                lambda name: name.replace(" ", "_")
-                .replace("-", "_")
-                .replace(".", "")
-                .replace("'", "_")
-                .strip()
-                .upper()
-            )
-            county_state = county_key.str.cat(state_code, sep="_")
-            county_names = pd.Series(
-                np.arange(len(County._member_names_)),
-                index=County._member_names_,
-            )
-            return county_names[county_state]
+            return map_county_string_to_enum(county_name, state_code)
 
         # Attempt to look up from ZIP code
         zip_code = household("zip_code", period).astype(int)
         zip_codes = ZIP_CODE_DATASET.set_index("zip_code")
         county_name = zip_codes.county[zip_code]
         state_code = zip_codes.state[zip_code]
-        county_key = county_name.apply(
-            lambda name: name.replace(" ", "_")
-            .replace("-", "_")
-            .replace(".", "")
-            .replace("'", "_")
-            .strip()
-            .upper()
-        )
-        county_state = county_key.str.cat(state_code, sep="_")
-        county_names = pd.Series(
-            np.arange(len(County._member_names_)), index=County._member_names_
-        )
-        return county_names[county_state]
+        return map_county_string_to_enum(county_name, state_code)