Skip to content
Merged
4 changes: 3 additions & 1 deletion _delphi_utils_python/delphi_utils/geomap.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,7 +305,9 @@ def add_geocode(
)

# state codes are all stored in one table
if new_code in state_codes:
if from_code in state_codes and new_code in state_codes:
crosswalk = self._load_crosswalk(from_code="state", to_code="state")
elif new_code in state_codes:
crosswalk = self._load_crosswalk(from_code=from_code, to_code="state")
crosswalk = crosswalk.rename(
columns={from_code: from_col, new_code: new_col}
Expand Down
5 changes: 5 additions & 0 deletions _delphi_utils_python/tests/test_geomap.py
Original file line number Diff line number Diff line change
Expand Up @@ -382,6 +382,11 @@ def test_add_geocode(self):
new_data2 = gmpr.add_geocode(new_data, "state_code", "hhs_region_number")
assert new_data2["hhs_region_number"].unique().size == 2

# state_name -> state_id
new_data = gmpr.add_geocode(self.zip_data, "zip", "state_name")
new_data2 = gmpr.add_geocode(new_data, "state_name", "state_id")
assert new_data2.shape == (12, 6)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should there also be an assert statement for new_data?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nah, new_data is just the test data for state_name to state_id, but it starts in zip form


# fips -> nation
new_data = gmpr.replace_geocode(self.fips_data_5, "fips", "nation")
assert new_data.equals(
Expand Down