Skip to content

Commit 8081679

Browse files
chekosclaude
andauthored
Complete feature parity: stubbed params, datasets, functions, and CLI (#299)
* Complete feature parity cleanup: stubbed params, datasets, functions, and CLI Implement all remaining items from the tidycensus evaluation checklist: - Stubbed parameters: time_series and breakdown_labels in estimates.py, breakdown_labels in flows.py, keep_geo_vars in get_acs/get_decennial - Datasets: pums_variables (API-fetched with cache), acs5_geography, mig_recodes - Functions: to_survey() with replicate-weight SE, interpolate_pw() for population-weighted areal interpolation, summary_files(), get_pop_groups() - CLI: config set-key, acs, decennial, variables, estimates commands - Tests: 70+ new test cases across 7 test files (162 total passing) - Eval doc: Section 11 checklist updated from Planned to Done Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix acs5_geography test mock path for Python 3.10 compatibility Use sys.modules to get the actual module object instead of import-as, which resolves to the re-exported function. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Fix bugs found in code review - Fix interpolate_pw() index mismatch when to_gdf has non-default index - Add validation: time_series=True only for product='population' - Rename to_survey() param type→weight_type to avoid shadowing builtin - Simplify get_pop_groups() state filter (avoid redundant JSON re-parsing) - Fix flow breakdown_labels zero-padded code mismatch (drop .astype(str)) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 1420d5e commit 8081679

21 files changed

Lines changed: 1562 additions & 32 deletions

TIDYCENSUS_EVALUATION.md

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -746,29 +746,29 @@ This is Pythonic and matches the pandas ecosystem.
746746

747747
| tidycensus Feature | pypums Equivalent | Phase | Status |
748748
|---|---|---|---|
749-
| `census_api_key()` | `census_api_key()` | 0 | Planned |
750-
| `get_acs()` | `get_acs()` | 1 | Planned |
751-
| `get_decennial()` | `get_decennial()` | 1 | Planned |
752-
| `load_variables()` | `load_variables()` | 1 | Planned |
753-
| `get_pums()` | `get_pums()` | 2 | Planned |
754-
| `get_estimates()` | `get_estimates()` | 3 | Planned |
755-
| `get_flows()` | `get_flows()` | 3 | Planned |
756-
| `moe_sum()` | `moe_sum()` | 2 | Planned |
757-
| `moe_prop()` | `moe_prop()` | 2 | Planned |
758-
| `moe_ratio()` | `moe_ratio()` | 2 | Planned |
759-
| `moe_product()` | `moe_product()` | 2 | Planned |
760-
| `significance()` | `significance()` | 2 | Planned |
761-
| `geometry=TRUE` | `geometry=True` | 2 | Planned |
762-
| `as_dot_density()` | `as_dot_density()` | 3 | Planned |
763-
| `interpolate_pw()` | `interpolate_pw()` | 3 | Planned |
764-
| `to_survey()` | `to_survey()` | 3 | Planned |
765-
| `fips_codes` | `fips_codes` | 0 | Planned |
766-
| `pums_variables` | `pums_variables` | 2 | Planned |
767-
| `mig_recodes` | `mig_recodes` | 3 | Planned |
749+
| `census_api_key()` | `census_api_key()` | 0 | Done |
750+
| `get_acs()` | `get_acs()` | 1 | Done |
751+
| `get_decennial()` | `get_decennial()` | 1 | Done |
752+
| `load_variables()` | `load_variables()` | 1 | Done |
753+
| `get_pums()` | `get_pums()` | 2 | Done |
754+
| `get_estimates()` | `get_estimates()` | 3 | Done |
755+
| `get_flows()` | `get_flows()` | 3 | Done |
756+
| `moe_sum()` | `moe_sum()` | 2 | Done |
757+
| `moe_prop()` | `moe_prop()` | 2 | Done |
758+
| `moe_ratio()` | `moe_ratio()` | 2 | Done |
759+
| `moe_product()` | `moe_product()` | 2 | Done |
760+
| `significance()` | `significance()` | 2 | Done |
761+
| `geometry=TRUE` | `geometry=True` | 2 | Done |
762+
| `as_dot_density()` | `as_dot_density()` | 3 | Done |
763+
| `interpolate_pw()` | `interpolate_pw()` | 3 | Done |
764+
| `to_survey()` | `to_survey()` | 3 | Done |
765+
| `fips_codes` | `fips_codes` | 0 | Done |
766+
| `pums_variables` | `pums_variables()` | 2 | Done |
767+
| `mig_recodes` | `mig_recodes` | 3 | Done |
768768
| `state_laea` / `county_laea` | deferred || Low priority |
769-
| `acs5_geography` | `acs5_geography` | 2 | Planned |
770-
| `summary_files()` | `summary_files()` | 3 | Planned |
771-
| `get_pop_groups()` | `get_pop_groups()` | 3 | Planned |
769+
| `acs5_geography` | `acs5_geography()` | 2 | Done |
770+
| `summary_files()` | `summary_files()` | 3 | Done |
771+
| `get_pop_groups()` | `get_pop_groups()` | 3 | Done |
772772
| `check_ddhca_groups()` | deferred || Low priority |
773773
| `shift_geo` (deprecated) | not implementing || Skip |
774774

pypums/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55

66
from .acs import get_acs as get_acs
77
from .api.key import census_api_key as census_api_key
8+
from .census_helpers import get_pop_groups as get_pop_groups
9+
from .census_helpers import summary_files as summary_files
810
from .constants import __app_name__ as __app_name__
911
from .constants import __version__ as __version__
1012
from .decennial import get_decennial as get_decennial
@@ -17,4 +19,5 @@
1719
from .moe import significance as significance
1820
from .pums import get_pums as get_pums
1921
from .survey import get_survey_metadata as get_survey_metadata
22+
from .survey import to_survey as to_survey
2023
from .variables import load_variables as load_variables

pypums/acs.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ def get_acs(
6363
moe_level: int = 90,
6464
summary_var: str | None = None,
6565
geometry: bool = False,
66+
keep_geo_vars: bool = False,
6667
cache_table: bool = False,
6768
key: str | None = None,
6869
) -> pd.DataFrame:
@@ -92,6 +93,9 @@ def get_acs(
9293
Variable ID to include as denominator columns.
9394
geometry
9495
If True, return a GeoDataFrame with shapes.
96+
keep_geo_vars
97+
If True, preserve the raw FIPS columns (state, county, tract,
98+
etc.) in the output alongside GEOID.
9599
cache_table
96100
If True, cache the API response locally to avoid redundant calls.
97101
key
@@ -177,12 +181,17 @@ def get_acs(
177181
scale_factor = _Z_SCORES[moe_level] / _Z_SCORES[90]
178182
df[moe_cols] = df[moe_cols] * scale_factor
179183

184+
# Determine which geo FIPS columns to keep.
185+
geo_cols_present = [c for c in _GEO_COL_ORDER if c in df.columns]
186+
extra_geo = geo_cols_present if keep_geo_vars else []
187+
180188
if output == "wide":
181-
keep_cols = ["GEOID", "NAME"] + estimate_cols + moe_cols
189+
keep_cols = ["GEOID", "NAME"] + extra_geo + estimate_cols + moe_cols
182190
result = df[[c for c in keep_cols if c in df.columns]]
183191
else:
184192
# Tidy format: melt estimate and MOE columns separately, then merge.
185193
id_cols = ["GEOID", "NAME"] if "GEOID" in df.columns else ["NAME"]
194+
id_cols = id_cols + extra_geo
186195

187196
# Exclude summary_var columns from the main melt.
188197
summary_est_col = f"{summary_var}E" if summary_var else None

pypums/census_helpers.py

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
"""Census helper functions for summary files and population groups.
2+
3+
Provides ``summary_files()`` and ``get_pop_groups()`` for discovering
4+
available Census Bureau data products.
5+
"""
6+
7+
import pandas as pd
8+
9+
from pypums.api.client import CENSUS_API_BASE, fetch_json
10+
11+
12+
def _fetch_json(url: str) -> dict | list:
13+
"""Thin wrapper so tests can mock ``pypums.census_helpers._fetch_json``."""
14+
return fetch_json(url)
15+
16+
17+
def summary_files(
18+
year: int = 2020,
19+
) -> pd.DataFrame:
20+
"""List available summary files for a given Census year.
21+
22+
Queries the Census API to discover which summary files (datasets)
23+
are available for a decennial census year.
24+
25+
Parameters
26+
----------
27+
year
28+
Census year (e.g. 2020, 2010).
29+
30+
Returns
31+
-------
32+
pd.DataFrame
33+
Columns: ``dataset_name``, ``title``, ``description``.
34+
"""
35+
url = f"{CENSUS_API_BASE}/{year}.json"
36+
raw = _fetch_json(url)
37+
38+
datasets = raw.get("dataset", []) if isinstance(raw, dict) else []
39+
40+
rows = []
41+
for ds in datasets:
42+
c_dataset = ds.get("c_dataset", [])
43+
# Only include decennial census datasets.
44+
if c_dataset and c_dataset[0].startswith("dec"):
45+
rows.append(
46+
{
47+
"dataset_name": "/".join(c_dataset),
48+
"title": ds.get("title", ""),
49+
"description": ds.get("description", ""),
50+
}
51+
)
52+
53+
return pd.DataFrame(rows, columns=["dataset_name", "title", "description"])
54+
55+
56+
def get_pop_groups(
57+
year: int = 2020,
58+
state: str | None = None,
59+
) -> pd.DataFrame:
60+
"""List available population groups for decennial Census DHC-A data.
61+
62+
Population groups are used with the ``pop_group`` parameter in
63+
``get_decennial()`` to access disaggregated data by detailed
64+
race/ethnicity categories.
65+
66+
Parameters
67+
----------
68+
year
69+
Census year (default 2020).
70+
state
71+
State FIPS code to filter population groups (optional).
72+
73+
Returns
74+
-------
75+
pd.DataFrame
76+
Columns: ``code``, ``label``.
77+
"""
78+
# The DHC-A population groups endpoint.
79+
url = f"{CENSUS_API_BASE}/{year}/dec/dhc-a/popgroup.json"
80+
raw = _fetch_json(url)
81+
82+
if not isinstance(raw, list) or len(raw) < 2:
83+
return pd.DataFrame(columns=["code", "label"])
84+
85+
# Census API returns JSON rows: first row is headers, rest is data.
86+
headers = raw[0]
87+
rows = []
88+
state_col = None
89+
if state is not None:
90+
state_col = next((h for h in headers if h.lower() == "state"), None)
91+
92+
for record in raw[1:]:
93+
row = dict(zip(headers, record, strict=True))
94+
rows.append(
95+
{
96+
"code": row.get("POP_GROUP", ""),
97+
"label": row.get("POP_GROUP_NAME", row.get("NAME", "")),
98+
"_state": row.get(state_col, "") if state_col else "",
99+
}
100+
)
101+
102+
df = pd.DataFrame(rows)
103+
104+
if state is not None and state_col is not None:
105+
df = df[df["_state"] == state].reset_index(drop=True)
106+
107+
return df.drop(columns=["_state"], errors="ignore")

pypums/cli.py

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,154 @@ def download_acs(
9696
_download_data(url, "acs", data_directory, extract)
9797

9898

99+
@cli.command("config")
100+
def config_set_key(
101+
key: str = typer.Argument(..., help="Census API key to store"),
102+
install: bool = typer.Option(
103+
True, "--install/--no-install", help="Save key to environment"
104+
),
105+
):
106+
"""Set your Census API key."""
107+
import os
108+
109+
from .api.key import census_api_key
110+
111+
os.environ["CENSUS_API_KEY"] = key
112+
census_api_key(key)
113+
console.print("[green]Census API key set successfully.[/green]")
114+
if install:
115+
console.print(
116+
"Key stored in environment variable CENSUS_API_KEY for this session."
117+
)
118+
119+
120+
@cli.command("acs")
121+
def acs_data(
122+
geography: str = typer.Argument(
123+
..., help="Geography level (e.g. 'state', 'county')"
124+
),
125+
variables: str = typer.Option(
126+
None, "--variables", "-v", help="Comma-separated variable IDs"
127+
),
128+
table: str = typer.Option(None, "--table", "-t", help="Census table ID"),
129+
state_opt: str = typer.Option(
130+
None, "--state", "-s", help="State FIPS or abbreviation"
131+
),
132+
county_opt: str = typer.Option(None, "--county", help="County FIPS code"),
133+
year_opt: int = typer.Option(2023, "--year", "-y", help="Data year"),
134+
survey_opt: str = typer.Option("acs5", "--survey", help="Survey: acs1 or acs5"),
135+
output: str = typer.Option("tidy", "--output", "-o", help="Output: tidy or wide"),
136+
key: str = typer.Option(None, "--key", "-k", help="Census API key"),
137+
):
138+
"""Fetch ACS data from the Census API."""
139+
from .acs import get_acs
140+
141+
var_list = variables.split(",") if variables else None
142+
df = get_acs(
143+
geography=geography,
144+
variables=var_list,
145+
table=table,
146+
state=state_opt,
147+
county=county_opt,
148+
year=year_opt,
149+
survey=survey_opt,
150+
output=output,
151+
key=key,
152+
)
153+
console.print(df.to_string())
154+
155+
156+
@cli.command("decennial")
157+
def decennial_data(
158+
geography: str = typer.Argument(
159+
..., help="Geography level (e.g. 'state', 'county')"
160+
),
161+
variables: str = typer.Option(
162+
None, "--variables", "-v", help="Comma-separated variable IDs"
163+
),
164+
table: str = typer.Option(None, "--table", "-t", help="Census table ID"),
165+
state_opt: str = typer.Option(
166+
None, "--state", "-s", help="State FIPS or abbreviation"
167+
),
168+
county_opt: str = typer.Option(None, "--county", help="County FIPS code"),
169+
year_opt: int = typer.Option(2020, "--year", "-y", help="Census year"),
170+
output: str = typer.Option("tidy", "--output", "-o", help="Output: tidy or wide"),
171+
key: str = typer.Option(None, "--key", "-k", help="Census API key"),
172+
):
173+
"""Fetch Decennial Census data from the Census API."""
174+
from .decennial import get_decennial
175+
176+
var_list = variables.split(",") if variables else None
177+
df = get_decennial(
178+
geography=geography,
179+
variables=var_list,
180+
table=table,
181+
state=state_opt,
182+
county=county_opt,
183+
year=year_opt,
184+
output=output,
185+
key=key,
186+
)
187+
console.print(df.to_string())
188+
189+
190+
@cli.command("variables")
191+
def variables_cmd(
192+
year_opt: int = typer.Option(2023, "--year", "-y", help="Data year"),
193+
dataset: str = typer.Option("acs5", "--dataset", "-d", help="Dataset identifier"),
194+
search: str = typer.Option(None, "--search", help="Filter by name/label/concept"),
195+
cache: bool = typer.Option(False, "--cache", help="Cache results"),
196+
):
197+
"""Search/browse Census variables."""
198+
from .variables import load_variables
199+
200+
df = load_variables(year=year_opt, dataset=dataset, cache=cache)
201+
if search:
202+
mask = (
203+
df["name"].str.contains(search, case=False, na=False)
204+
| df["label"].str.contains(search, case=False, na=False)
205+
| df["concept"].str.contains(search, case=False, na=False)
206+
)
207+
df = df[mask]
208+
console.print(df.to_string())
209+
210+
211+
@cli.command("estimates")
212+
def estimates_cmd(
213+
geography: str = typer.Argument(
214+
..., help="Geography level (e.g. 'state', 'county')"
215+
),
216+
product: str = typer.Option(
217+
"population", "--product", "-p", help="Estimates product"
218+
),
219+
variables: str = typer.Option(
220+
None, "--variables", "-v", help="Comma-separated variable IDs"
221+
),
222+
state_opt: str = typer.Option(
223+
None, "--state", "-s", help="State FIPS or abbreviation"
224+
),
225+
county_opt: str = typer.Option(None, "--county", help="County FIPS code"),
226+
vintage: int = typer.Option(2023, "--vintage", help="Vintage year"),
227+
output: str = typer.Option("tidy", "--output", "-o", help="Output: tidy or wide"),
228+
key: str = typer.Option(None, "--key", "-k", help="Census API key"),
229+
):
230+
"""Fetch population estimates from the Census API."""
231+
from .estimates import get_estimates
232+
233+
var_list = variables.split(",") if variables else None
234+
df = get_estimates(
235+
geography=geography,
236+
product=product,
237+
variables=var_list,
238+
state=state_opt,
239+
county=county_opt,
240+
vintage=vintage,
241+
output=output,
242+
key=key,
243+
)
244+
console.print(df.to_string())
245+
246+
99247
def _version_callback(value: bool) -> None:
100248
if value:
101249
typer.echo(f"{__app_name__} v{__version__}")

pypums/datasets/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
"""Built-in reference datasets."""
22

3+
from pypums.datasets.acs5_geography import acs5_geography as acs5_geography
34
from pypums.datasets.fips import fips_codes as fips_codes
45
from pypums.datasets.fips import lookup_fips as lookup_fips
56
from pypums.datasets.fips import lookup_name as lookup_name
7+
from pypums.datasets.mig_recodes import mig_recodes as mig_recodes
8+
from pypums.datasets.pums_vars import pums_variables as pums_variables

0 commit comments

Comments
 (0)