diff --git a/src/geocodio/client.py b/src/geocodio/client.py index dbf3c04..b74cf2e 100644 --- a/src/geocodio/client.py +++ b/src/geocodio/client.py @@ -443,20 +443,17 @@ def _parse_fields(self, fields_data: dict | None) -> GeocodioFields | None: for district in fields_data["school"] ] - census2010 = ( - CensusData.from_api(fields_data["census2010"]) - if "census2010" in fields_data else None - ) - - census2020 = ( - CensusData.from_api(fields_data["census2020"]) - if "census2020" in fields_data else None - ) - - census2023 = ( - CensusData.from_api(fields_data["census2023"]) - if "census2023" in fields_data else None - ) + # Dynamically parse all census fields (e.g., census2010, census2020, census2024, etc.) + # This supports any census year returned by the API + from dataclasses import fields as dataclass_fields + valid_field_names = {f.name for f in dataclass_fields(GeocodioFields)} + + census_fields = {} + for key in fields_data: + if key.startswith("census") and key[6:].isdigit(): # e.g., "census2024" + # Only include if it's a defined field in GeocodioFields + if key in valid_field_names: + census_fields[key] = CensusData.from_api(fields_data[key]) acs = ( ACSSurveyData.from_api(fields_data["acs"]) @@ -515,9 +512,6 @@ def _parse_fields(self, fields_data: dict | None) -> GeocodioFields | None: state_legislative_districts=state_legislative_districts, state_legislative_districts_next=state_legislative_districts_next, school_districts=school_districts, - census2010=census2010, - census2020=census2020, - census2023=census2023, acs=acs, demographics=demographics, economics=economics, @@ -528,6 +522,7 @@ def _parse_fields(self, fields_data: dict | None) -> GeocodioFields | None: provriding=provriding, provriding_next=provriding_next, statcan=statcan, + **census_fields, # Dynamically include all census year fields ) # @TODO add a "keep_trying" parameter to download() to keep trying until the list is processed. diff --git a/src/geocodio/models.py b/src/geocodio/models.py index ec4f906..32ad22b 100644 --- a/src/geocodio/models.py +++ b/src/geocodio/models.py @@ -283,7 +283,7 @@ class FFIECData(ApiModelMixin): extras: Dict[str, Any] = field(default_factory=dict, repr=False) -@dataclass(slots=True, frozen=True) +@dataclass(frozen=True) class GeocodioFields: """ Container for optional 'fields' returned by the Geocodio API. diff --git a/tests/unit/test_geocode.py b/tests/unit/test_geocode.py index 77c5365..6b0a710 100644 --- a/tests/unit/test_geocode.py +++ b/tests/unit/test_geocode.py @@ -498,4 +498,102 @@ def batch_response_callback(request): assert resp.results[1].formatted_address == "638 E 13th Ave, Denver, CO 80203" assert resp.results[1].fields.timezone.name == "America/Denver" assert resp.results[1].fields.timezone.utc_offset == -7 - assert resp.results[1].fields.congressional_districts[0].district_number == 1 \ No newline at end of file + assert resp.results[1].fields.congressional_districts[0].district_number == 1 + + +def test_geocode_with_census_fields(client, httpx_mock): + """Test geocoding with census field appends including all census years.""" + # Arrange: stub the API call with multiple census years + def response_callback(request): + assert request.method == "GET" + assert request.url.params["fields"] == "census2010,census2020,census2023,census2024" + return httpx.Response(200, json={ + "results": [{ + "address_components": { + "number": "1640", + "street": "Main", + "suffix": "St", + "city": "Sheldon", + "state": "VT", + "zip": "05483", + "country": "US" + }, + "formatted_address": "1640 Main St, Sheldon, VT 05483", + "location": {"lat": 44.895469, "lng": -72.953264}, + "accuracy": 1, + "accuracy_type": "rooftop", + "source": "Vermont", + "fields": { + "census2010": { + "tract": "960100", + "block": "2001", + "blockgroup": "2", + "county_fips": "50011", + "state_fips": "50" + }, + "census2020": { + "tract": "960100", + "block": "2002", + "blockgroup": "2", + "county_fips": "50011", + "state_fips": "50" + }, + "census2023": { + "tract": "960100", + "block": "2003", + "blockgroup": "2", + "county_fips": "50011", + "state_fips": "50" + }, + "census2024": { + "tract": "960100", + "block": "2004", + "blockgroup": "2", + "county_fips": "50011", + "state_fips": "50" + } + } + }] + }) + + httpx_mock.add_callback( + callback=response_callback, + url=httpx.URL("https://api.test/v1.9/geocode", params={ + "street": "1640 Main St", + "city": "Sheldon", + "state": "VT", + "postal_code": "05483", + "fields": "census2010,census2020,census2023,census2024" + }), + match_headers={"Authorization": "Bearer TEST_KEY"}, + ) + + # Act + resp = client.geocode( + {"city": "Sheldon", "state": "VT", "street": "1640 Main St", "postal_code": "05483"}, + fields=["census2010", "census2020", "census2023", "census2024"], + ) + + # Assert + assert len(resp.results) == 1 + result = resp.results[0] + assert result.formatted_address == "1640 Main St, Sheldon, VT 05483" + + # Check that all census fields are present and parsed correctly + assert result.fields.census2010 is not None + assert result.fields.census2010.tract == "960100" + assert result.fields.census2010.block == "2001" + assert result.fields.census2010.county_fips == "50011" + + assert result.fields.census2020 is not None + assert result.fields.census2020.tract == "960100" + assert result.fields.census2020.block == "2002" + + assert result.fields.census2023 is not None + assert result.fields.census2023.tract == "960100" + assert result.fields.census2023.block == "2003" + + # This will fail until we fix the parsing logic + assert result.fields.census2024 is not None + assert result.fields.census2024.tract == "960100" + assert result.fields.census2024.block == "2004" \ No newline at end of file