Skip to content

Commit 4925b54

Browse files
chekosclaude
chekos
and
claude
authored
Add Phase 1 core data functions: get_acs, get_decennial, load_variables (#295)
* Add Phase 1 core data functions: get_acs, get_decennial, load_variables Implement the three core Census data retrieval functions: - get_acs(): ACS data with tidy/wide output, MOE scaling (90/95/99%), summary variable support, and variable/table input modes - get_decennial(): Decennial Census data with tidy/wide output, automatic dataset selection by year (dhc/sf1), pop_group support - load_variables(): Variable metadata discovery with in-memory caching, supports all dataset types (acs5, acs1, pl, subject, profile) All three use Phase 0 infrastructure (census_api_key, build_geography_query). Internal _call_census_api functions are mockable for testing. All 24 Phase 1 tests pass. All 48 prior tests unaffected. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Address PR review feedback - Raise NotImplementedError when geometry=True (not yet supported) - Validate output param (must be 'tidy' or 'wide') and moe_level param (must be 90, 95, or 99) with clear error messages - Raise ValueError for unsupported decennial years instead of silently falling back to wrong dataset - Extract shared _call_census_api and CENSUS_API_BASE into pypums/api/client.py to eliminate duplication - Replace iterrows() loops with vectorized pd.melt() for tidy conversion in both acs.py and decennial.py Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix pop_group year validation bypass, strengthen dataset routing tests - Validate year before checking pop_group so unsupported years always raise ValueError regardless of pop_group - Strengthen test assertions for dataset routing: verify the URL contains "dec/dhc" for 2020 and "dec/sf1" for 2010, instead of only asserting call_args is not None Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Consolidate HTTP logic into shared client - Add fetch_json() to pypums/api/client.py for simple GET requests - Extract CENSUS_TIMEOUT constant (single source of truth) - Remove direct httpx usage from variables.py, delegate to client Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 58615e1 commit 4925b54

6 files changed

Lines changed: 411 additions & 4 deletions

File tree

‎pypums/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
# type: ignore[attr-defined]
22
"""Download PUMS data files from the US Census Bureau's FTP server."""
33

4+
from pypums.acs import get_acs as get_acs
45
from pypums.api.key import census_api_key as census_api_key
6+
from pypums.decennial import get_decennial as get_decennial
57
from pypums.surveys import ACS as ACS
8+
from pypums.variables import load_variables as load_variables
69

710
from .constants import __app_name__ as __app_name__
811
from .constants import __version__ as __version__

‎pypums/acs.py‎

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
"""American Community Survey data retrieval via the Census API."""
2+
3+
import pandas as pd
4+
from pypums.api.client import CENSUS_API_BASE, call_census_api
5+
from pypums.api.geography import build_geography_query
6+
from pypums.api.key import census_api_key
7+
8+
# Z-scores for MOE confidence levels.
9+
_Z_SCORES: dict[int, float] = {
10+
90: 1.645,
11+
95: 1.960,
12+
99: 2.576,
13+
}
14+
15+
# Geography columns that appear in Census API responses.
16+
_GEO_COLUMNS = frozenset({
17+
"state", "county", "tract", "block group", "block",
18+
"place", "congressional district",
19+
"state legislative district (upper chamber)",
20+
"state legislative district (lower chamber)",
21+
"zip code tabulation area",
22+
"school district (unified)",
23+
"school district (elementary)",
24+
"school district (secondary)",
25+
"metropolitan statistical area/micropolitan statistical area",
26+
"combined statistical area",
27+
"public use microdata area",
28+
"american indian area/alaska native area/hawaiian home land",
29+
"us", "region", "division", "county subdivision",
30+
})
31+
32+
33+
def _call_census_api(url: str, params: dict) -> list[list[str]]:
34+
"""Thin wrapper so tests can mock ``pypums.acs._call_census_api``."""
35+
return call_census_api(url, params)
36+
37+
38+
def get_acs(
39+
geography: str,
40+
variables: str | list[str] | None = None,
41+
table: str | None = None,
42+
state: str | None = None,
43+
county: str | None = None,
44+
year: int = 2023,
45+
output: str = "tidy",
46+
moe_level: int = 90,
47+
summary_var: str | None = None,
48+
geometry: bool = False,
49+
key: str | None = None,
50+
) -> pd.DataFrame:
51+
"""Retrieve American Community Survey data from the Census API.
52+
53+
Parameters
54+
----------
55+
geography
56+
Geography level (e.g. ``"state"``, ``"county"``, ``"tract"``).
57+
variables
58+
Variable ID or list of IDs (e.g. ``"B01001_001"``).
59+
table
60+
Census table ID (e.g. ``"B01001"``). Alternative to *variables*.
61+
state
62+
State FIPS code or abbreviation.
63+
county
64+
County FIPS code.
65+
year
66+
Data year (default 2023).
67+
output
68+
``"tidy"`` (default) or ``"wide"``.
69+
moe_level
70+
Confidence level for MOE: 90, 95, or 99 (default 90).
71+
summary_var
72+
Variable ID to include as denominator columns.
73+
geometry
74+
If True, return a GeoDataFrame with shapes (not yet implemented).
75+
key
76+
Census API key. Falls back to ``census_api_key()``.
77+
78+
Returns
79+
-------
80+
pd.DataFrame
81+
Census data in tidy or wide format.
82+
"""
83+
if geometry:
84+
raise NotImplementedError("geometry=True is not yet supported.")
85+
if output not in ("tidy", "wide"):
86+
raise ValueError(f"output must be 'tidy' or 'wide', got {output!r}")
87+
if moe_level not in _Z_SCORES:
88+
raise ValueError(f"moe_level must be one of {sorted(_Z_SCORES)}, got {moe_level!r}")
89+
90+
api_key = census_api_key(key) if key else census_api_key()
91+
for_clause, in_clause = build_geography_query(geography, state=state, county=county)
92+
93+
# Build the variable list for the API request.
94+
if variables is not None:
95+
if isinstance(variables, str):
96+
variables = [variables]
97+
# Census API needs E/M suffixes for ACS.
98+
api_vars = []
99+
for v in variables:
100+
api_vars.append(f"{v}E")
101+
api_vars.append(f"{v}M")
102+
elif table is not None:
103+
api_vars = [f"group({table})"]
104+
else:
105+
raise ValueError("Must provide either 'variables' or 'table'.")
106+
107+
# Add summary variable if requested.
108+
if summary_var is not None:
109+
api_vars.append(f"{summary_var}E")
110+
api_vars.append(f"{summary_var}M")
111+
112+
url = f"{CENSUS_API_BASE}/{year}/acs/acs5"
113+
params: dict[str, str] = {
114+
"get": f"NAME,{','.join(api_vars)}",
115+
"for": for_clause,
116+
"key": api_key,
117+
}
118+
if in_clause is not None:
119+
params["in"] = in_clause
120+
121+
data = _call_census_api(url, params)
122+
123+
# Convert JSON rows to DataFrame.
124+
headers = data[0]
125+
df = pd.DataFrame(data[1:], columns=headers)
126+
127+
# Build GEOID from FIPS columns.
128+
geo_cols = [c for c in df.columns if c in _GEO_COLUMNS]
129+
if geo_cols:
130+
df["GEOID"] = df[geo_cols].apply(lambda row: "".join(row), axis=1)
131+
132+
# Identify estimate and MOE columns.
133+
estimate_cols = [c for c in df.columns if c.endswith("E") and c != "NAME"]
134+
moe_cols = [c for c in df.columns if c.endswith("M")]
135+
136+
# Convert to numeric.
137+
for col in estimate_cols + moe_cols:
138+
df[col] = pd.to_numeric(df[col], errors="coerce")
139+
140+
# Scale MOE if needed.
141+
if moe_level != 90:
142+
scale_factor = _Z_SCORES[moe_level] / _Z_SCORES[90]
143+
df[moe_cols] = df[moe_cols] * scale_factor
144+
145+
if output == "wide":
146+
keep_cols = ["GEOID", "NAME"] + estimate_cols + moe_cols
147+
return df[[c for c in keep_cols if c in df.columns]]
148+
149+
# Tidy format: melt estimate and MOE columns separately, then merge.
150+
id_cols = ["GEOID", "NAME"] if "GEOID" in df.columns else ["NAME"]
151+
152+
# Exclude summary_var columns from the main melt.
153+
summary_est_col = f"{summary_var}E" if summary_var else None
154+
summary_moe_col = f"{summary_var}M" if summary_var else None
155+
main_est_cols = [c for c in estimate_cols if c != summary_est_col]
156+
main_moe_cols = [c for c in moe_cols if c != summary_moe_col]
157+
158+
est_long = df.melt(
159+
id_vars=id_cols,
160+
value_vars=main_est_cols,
161+
var_name="_est_var",
162+
value_name="estimate",
163+
)
164+
est_long["variable"] = est_long["_est_var"].str[:-1]
165+
166+
moe_long = df.melt(
167+
id_vars=id_cols,
168+
value_vars=main_moe_cols,
169+
var_name="_moe_var",
170+
value_name="moe",
171+
)
172+
moe_long["variable"] = moe_long["_moe_var"].str[:-1]
173+
174+
tidy = est_long[id_cols + ["variable", "estimate"]].merge(
175+
moe_long[id_cols + ["variable", "moe"]],
176+
on=id_cols + ["variable"],
177+
)
178+
179+
# Add summary variable columns if requested.
180+
if summary_var is not None and summary_est_col in df.columns:
181+
summary_df = df[id_cols + [summary_est_col, summary_moe_col]].rename(
182+
columns={summary_est_col: "summary_est", summary_moe_col: "summary_moe"},
183+
)
184+
tidy = tidy.merge(summary_df, on=id_cols)
185+
186+
return tidy

‎pypums/api/client.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
"""Shared Census API HTTP client."""
2+
3+
import httpx
4+
5+
CENSUS_API_BASE = "https://api.census.gov/data"
6+
CENSUS_TIMEOUT = 30
7+
8+
9+
def call_census_api(url: str, params: dict) -> list[list[str]]:
10+
"""Make an HTTP request to the Census API and return JSON rows."""
11+
response = httpx.get(url, params=params, timeout=CENSUS_TIMEOUT)
12+
response.raise_for_status()
13+
return response.json()
14+
15+
16+
def fetch_json(url: str) -> dict:
17+
"""Fetch JSON from a Census API endpoint."""
18+
response = httpx.get(url, timeout=CENSUS_TIMEOUT)
19+
response.raise_for_status()
20+
return response.json()

‎pypums/decennial.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
"""Decennial Census data retrieval via the Census API."""
2+
3+
import pandas as pd
4+
from pypums.api.client import CENSUS_API_BASE, call_census_api
5+
from pypums.api.geography import build_geography_query
6+
from pypums.api.key import census_api_key
7+
8+
# Default dataset by census year.
9+
_YEAR_DATASETS: dict[int, str] = {
10+
2020: "dec/dhc",
11+
2010: "dec/sf1",
12+
2000: "dec/sf1",
13+
}
14+
15+
# Geography columns that appear in Census API responses.
16+
_GEO_COLUMNS = frozenset({
17+
"state", "county", "tract", "block group", "block",
18+
"place", "congressional district",
19+
"us", "region", "division", "county subdivision",
20+
})
21+
22+
23+
def _call_census_api(url: str, params: dict) -> list[list[str]]:
24+
"""Thin wrapper so tests can mock ``pypums.decennial._call_census_api``."""
25+
return call_census_api(url, params)
26+
27+
28+
def get_decennial(
29+
geography: str,
30+
variables: str | list[str] | None = None,
31+
table: str | None = None,
32+
state: str | None = None,
33+
county: str | None = None,
34+
year: int = 2020,
35+
output: str = "tidy",
36+
pop_group: str | None = None,
37+
geometry: bool = False,
38+
key: str | None = None,
39+
) -> pd.DataFrame:
40+
"""Retrieve Decennial Census data from the Census API.
41+
42+
Parameters
43+
----------
44+
geography
45+
Geography level (e.g. ``"state"``, ``"county"``).
46+
variables
47+
Variable ID or list of IDs (e.g. ``"P1_001N"``).
48+
table
49+
Census table ID. Alternative to *variables*.
50+
state
51+
State FIPS code or abbreviation.
52+
county
53+
County FIPS code.
54+
year
55+
Census year: 2000, 2010, or 2020 (default 2020).
56+
output
57+
``"tidy"`` (default) or ``"wide"``.
58+
pop_group
59+
Population group code for DHC-A disaggregated data.
60+
geometry
61+
If True, return a GeoDataFrame with shapes (not yet implemented).
62+
key
63+
Census API key. Falls back to ``census_api_key()``.
64+
65+
Returns
66+
-------
67+
pd.DataFrame
68+
Census data in tidy or wide format.
69+
"""
70+
if geometry:
71+
raise NotImplementedError("geometry=True is not yet supported.")
72+
if output not in ("tidy", "wide"):
73+
raise ValueError(f"output must be 'tidy' or 'wide', got {output!r}")
74+
75+
api_key = census_api_key(key) if key else census_api_key()
76+
for_clause, in_clause = build_geography_query(geography, state=state, county=county)
77+
78+
# Validate year and select dataset.
79+
if year not in _YEAR_DATASETS:
80+
raise ValueError(
81+
f"Unsupported decennial year: {year}. "
82+
f"Supported years: {sorted(_YEAR_DATASETS)}"
83+
)
84+
dataset = "dec/dhc-a" if pop_group is not None else _YEAR_DATASETS[year]
85+
86+
# Build the variable list.
87+
if variables is not None:
88+
if isinstance(variables, str):
89+
variables = [variables]
90+
api_vars = list(variables)
91+
elif table is not None:
92+
api_vars = [f"group({table})"]
93+
else:
94+
raise ValueError("Must provide either 'variables' or 'table'.")
95+
96+
url = f"{CENSUS_API_BASE}/{year}/{dataset}"
97+
params: dict[str, str] = {
98+
"get": f"NAME,{','.join(api_vars)}",
99+
"for": for_clause,
100+
"key": api_key,
101+
}
102+
if in_clause is not None:
103+
params["in"] = in_clause
104+
if pop_group is not None:
105+
params["POP_GROUP"] = pop_group
106+
107+
data = _call_census_api(url, params)
108+
109+
# Convert JSON rows to DataFrame.
110+
headers = data[0]
111+
df = pd.DataFrame(data[1:], columns=headers)
112+
113+
# Build GEOID from FIPS columns.
114+
geo_cols = [c for c in df.columns if c in _GEO_COLUMNS]
115+
if geo_cols:
116+
df["GEOID"] = df[geo_cols].apply(lambda row: "".join(row), axis=1)
117+
118+
# Identify variable columns (everything except NAME and geo columns).
119+
var_cols = [c for c in df.columns if c not in _GEO_COLUMNS and c not in ("NAME", "GEOID")]
120+
121+
# Convert to numeric.
122+
for col in var_cols:
123+
df[col] = pd.to_numeric(df[col], errors="coerce")
124+
125+
if output == "wide":
126+
keep_cols = ["GEOID", "NAME"] + var_cols
127+
return df[[c for c in keep_cols if c in df.columns]]
128+
129+
# Tidy format: melt to one row per geography × variable.
130+
id_cols = ["GEOID", "NAME"] if "GEOID" in df.columns else ["NAME"]
131+
tidy = df.melt(
132+
id_vars=id_cols,
133+
value_vars=var_cols,
134+
var_name="variable",
135+
value_name="value",
136+
)
137+
138+
return tidy

0 commit comments

Comments
 (0)