refactor: Extract geo+vectors fetch and merge into dedicated functions

Extract the complex hybrid geo+vectors logic (~70 lines) from get_census()
into two dedicated helper functions:

1. _fetch_census_with_geometry_and_vectors():
   - Handles the CensusMapper API quirk where geo.geojson doesn't return
     vector data properly
   - Makes separate calls to geo.geojson and data.csv endpoints
   - Orchestrates the merge of geometry and vector data

2. _merge_geo_and_csv_results():
   - Merges GeoDataFrame with CSV DataFrame on geographic identifier
   - Detects common merge keys (GeoUID, id, rgid) automatically
   - Falls back to index-based merge if no common key found
   - Handles duplicate key cleanup after merge

This improves:
- Readability: get_census() main flow is now much simpler
- Testability: merge logic can be unit tested in isolation
- Single responsibility: each function has one clear purpose
- Documentation: detailed docstrings explain the API quirk

Added 2 new unit tests:
- test_merge_on_geoid_key: verifies key-based merge
- test_merge_fallback_by_index: verifies index fallback

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: dshkol

pycancensus/core.py

@@ -171,80 +171,10 @@ def get_census(
         # Handle geo_format='geopandas' with vectors using hybrid approach
         if geo_format == "geopandas" and vectors:
             # The geo.geojson endpoint doesn't properly return vector data
-            # So we need to fetch geometry and data separately, then merge
-
-            # 1. Fetch geometry data
-            geo_request_data = request_data.copy()
-            if "vectors" in geo_request_data:
-                del geo_request_data["vectors"]  # Remove vectors for geo request
-            if resolution == "high":
-                geo_request_data["resolution"] = "high"
-
-            geo_multipart_data = {}
-            for key, value in geo_request_data.items():
-                geo_multipart_data[key] = (None, value)
-
-            geo_response = get_session().post(
-                f"{base_url}geo.geojson", files=geo_multipart_data
+            # Use dedicated function to fetch and merge geo + vector data
+            result = _fetch_census_with_geometry_and_vectors(
+                base_url, request_data, resolution, vectors, labels
             )
-            geo_data = geo_response.json()
-            geo_result = _process_geojson_response(geo_data, None, labels)  # No vectors
-
-            # 2. Fetch vector data using CSV endpoint
-            csv_multipart_data = {}
-            for key, value in request_data.items():
-                csv_multipart_data[key] = (None, value)
-
-            csv_response = get_session().post(
-                f"{base_url}data.csv", files=csv_multipart_data
-            )
-            csv_result = _process_csv_response(csv_response.text, vectors, labels)
-
-            # 3. Merge the results
-            # Use a common identifier to merge - typically 'GeoUID' from CSV and 'id' from GeoJSON
-            merge_key_csv = None
-            merge_key_geo = None
-
-            # Find the appropriate merge keys
-            for potential_key in ["GeoUID", "id", "rgid"]:
-                if potential_key in csv_result.columns:
-                    merge_key_csv = potential_key
-                    break
-
-            for potential_key in ["id", "rgid", "GeoUID"]:
-                if potential_key in geo_result.columns:
-                    merge_key_geo = potential_key
-                    break
-
-            if merge_key_csv and merge_key_geo:
-                # Merge on the identifier
-                # Keep all columns from geo_result, add vector columns from csv_result
-                vector_columns = [
-                    col for col in csv_result.columns if col.startswith("v_")
-                ]
-                merge_columns = [merge_key_csv] + vector_columns
-
-                result = geo_result.merge(
-                    csv_result[merge_columns],
-                    left_on=merge_key_geo,
-                    right_on=merge_key_csv,
-                    how="left",
-                )
-
-                # Drop the duplicate merge key if it was added
-                if merge_key_csv != merge_key_geo and merge_key_csv in result.columns:
-                    result = result.drop(columns=[merge_key_csv])
-
-            else:
-                # Fallback: assume same order and merge by index
-                vector_columns = [
-                    col for col in csv_result.columns if col.startswith("v_")
-                ]
-                for col in vector_columns:
-                    if len(csv_result) == len(geo_result):
-                        geo_result[col] = csv_result[col].values
-                result = geo_result
-
         else:
             # Standard single-endpoint approach
             if geo_format == "geopandas":
@@ -308,6 +238,124 @@ def _generate_cache_key(dataset, regions, vectors, level, geo_format):
     return hashlib.md5(params_str.encode()).hexdigest()
 
 
+def _fetch_census_with_geometry_and_vectors(
+    base_url: str,
+    request_data: dict,
+    resolution: str,
+    vectors: List[str],
+    labels: str,
+) -> gpd.GeoDataFrame:
+    """
+    Fetch census data with both geometry and vector data.
+
+    The CensusMapper geo.geojson endpoint doesn't properly return vector data,
+    so this function makes separate calls to geo.geojson and data.csv endpoints,
+    then merges the results on geographic identifier.
+
+    Parameters
+    ----------
+    base_url : str
+        The API base URL (e.g., "https://censusmapper.ca/api/v1/").
+    request_data : dict
+        The base request parameters (dataset, level, api_key, regions, etc.).
+    resolution : str
+        Resolution of geographic data - 'simplified' or 'high'.
+    vectors : list of str
+        Vector codes to retrieve.
+    labels : str
+        Label format - 'detailed' or 'short'.
+
+    Returns
+    -------
+    gpd.GeoDataFrame
+        GeoDataFrame with geometry and vector data merged.
+    """
+    # 1. Fetch geometry data (without vectors)
+    geo_request_data = request_data.copy()
+    if "vectors" in geo_request_data:
+        del geo_request_data["vectors"]
+    if resolution == "high":
+        geo_request_data["resolution"] = "high"
+
+    geo_multipart_data = {key: (None, value) for key, value in geo_request_data.items()}
+    geo_response = get_session().post(
+        f"{base_url}geo.geojson", files=geo_multipart_data
+    )
+    geo_data = geo_response.json()
+    geo_result = _process_geojson_response(geo_data, None, labels)
+
+    # 2. Fetch vector data using CSV endpoint
+    csv_multipart_data = {key: (None, value) for key, value in request_data.items()}
+    csv_response = get_session().post(f"{base_url}data.csv", files=csv_multipart_data)
+    csv_result = _process_csv_response(csv_response.text, vectors, labels)
+
+    # 3. Merge the results on geographic identifier
+    return _merge_geo_and_csv_results(geo_result, csv_result)
+
+
+def _merge_geo_and_csv_results(
+    geo_result: gpd.GeoDataFrame,
+    csv_result: pd.DataFrame,
+) -> gpd.GeoDataFrame:
+    """
+    Merge GeoDataFrame with CSV DataFrame on geographic identifier.
+
+    Finds a common identifier column (GeoUID, id, or rgid) and merges
+    the vector columns from CSV onto the GeoDataFrame.
+
+    Parameters
+    ----------
+    geo_result : gpd.GeoDataFrame
+        GeoDataFrame with geometry data.
+    csv_result : pd.DataFrame
+        DataFrame with vector data.
+
+    Returns
+    -------
+    gpd.GeoDataFrame
+        Merged GeoDataFrame with geometry and vector columns.
+    """
+    # Find merge keys in each DataFrame
+    merge_key_csv = None
+    merge_key_geo = None
+
+    for potential_key in ["GeoUID", "id", "rgid"]:
+        if potential_key in csv_result.columns:
+            merge_key_csv = potential_key
+            break
+
+    for potential_key in ["id", "rgid", "GeoUID"]:
+        if potential_key in geo_result.columns:
+            merge_key_geo = potential_key
+            break
+
+    if merge_key_csv and merge_key_geo:
+        # Merge on identifier - keep geo columns, add vector columns from CSV
+        vector_columns = [col for col in csv_result.columns if col.startswith("v_")]
+        merge_columns = [merge_key_csv] + vector_columns
+
+        result = geo_result.merge(
+            csv_result[merge_columns],
+            left_on=merge_key_geo,
+            right_on=merge_key_csv,
+            how="left",
+        )
+
+        # Drop duplicate merge key if names differ
+        if merge_key_csv != merge_key_geo and merge_key_csv in result.columns:
+            result = result.drop(columns=[merge_key_csv])
+
+    else:
+        # Fallback: assume same row order and merge by index
+        vector_columns = [col for col in csv_result.columns if col.startswith("v_")]
+        result = geo_result.copy()
+        for col in vector_columns:
+            if len(csv_result) == len(geo_result):
+                result[col] = csv_result[col].values
+
+    return result
+
+
 def _extract_vector_metadata(df, vectors, labels):
     """Extract vector metadata from column names and store as attribute."""
     if not vectors:

tests/test_basic.py

@@ -359,6 +359,68 @@ def test_normalize_produces_equivalent_results(self):
         assert pd.isna(geo_result["pop"].iloc[1])
 
 
+class TestGeoVectorsMerge:
+    """Test geo+vectors merge functionality."""
+
+    def test_merge_on_geoid_key(self):
+        """Test merging geo and CSV results on GeoUID/id key."""
+        from pycancensus.core import _merge_geo_and_csv_results
+
+        # Create mock GeoDataFrame
+        geo_result = gpd.GeoDataFrame(
+            {
+                "id": ["001", "002", "003"],
+                "name": ["Region A", "Region B", "Region C"],
+                "geometry": [None, None, None],
+            }
+        )
+
+        # Create mock CSV result
+        csv_result = pd.DataFrame(
+            {
+                "GeoUID": ["001", "002", "003"],
+                "v_CA21_1": [100, 200, 300],
+                "v_CA21_2": [50, 60, 70],
+            }
+        )
+
+        result = _merge_geo_and_csv_results(geo_result, csv_result)
+
+        # Should have vector columns merged
+        assert "v_CA21_1" in result.columns
+        assert "v_CA21_2" in result.columns
+        assert list(result["v_CA21_1"]) == [100, 200, 300]
+
+        # Should have geo columns preserved
+        assert "name" in result.columns
+        assert "geometry" in result.columns
+
+    def test_merge_fallback_by_index(self):
+        """Test fallback merge by index when no common key found."""
+        from pycancensus.core import _merge_geo_and_csv_results
+
+        # Create mock data without common keys
+        geo_result = gpd.GeoDataFrame(
+            {
+                "custom_id": ["A", "B"],
+                "geometry": [None, None],
+            }
+        )
+
+        csv_result = pd.DataFrame(
+            {
+                "other_id": ["X", "Y"],
+                "v_CA21_1": [100, 200],
+            }
+        )
+
+        result = _merge_geo_and_csv_results(geo_result, csv_result)
+
+        # Should still merge by index
+        assert "v_CA21_1" in result.columns
+        assert list(result["v_CA21_1"]) == [100, 200]
+
+
 class TestCache:
     """Test caching functionality."""