Add Phase 3: estimates, flows, FIPS lookups, and caching integration

.github/PULL_REQUEST_TEMPLATE.md

@@ -24,4 +24,3 @@
 - [ ] I've read the [`CODE_OF_CONDUCT.md`](https://github.com/chekos/pypums/blob/master/CODE_OF_CONDUCT.md) document.
 - [ ] I've read the [`CONTRIBUTING.md`](https://github.com/chekos/pypums/blob/master/CONTRIBUTING.md) guide.
 - [ ] I've written tests for all new methods and classes that I created.
-

.github/workflows/claude.yml

@@ -47,4 +47,3 @@ jobs:
           # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
           # or https://code.claude.com/docs/en/cli-reference for available options
           # claude_args: '--allowed-tools Bash(gh pr:*)'
-

.pre-commit-config.yaml

@@ -5,6 +5,7 @@ repos:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: check-yaml
+        args: [--unsafe]
       - id: check-toml
       - id: check-added-large-files
   - repo: https://github.com/astral-sh/ruff-pre-commit

docs/about/contributing.md

@@ -73,7 +73,7 @@ local development.
     Or if you are using `pipenv`:
 
         pipenv shell
-    
+
     Now install the dependencies and test dependencies:
 
         pip install -e '.[test]'
@@ -82,7 +82,7 @@ local development.
 
         git checkout -b name-of-your-bugfix-or-feature
 
-    Now you can make your changes locally.  
+    Now you can make your changes locally.
     <br>
 
 5.  When you're done making changes, check that your changes pass
@@ -111,4 +111,3 @@ Before you submit a pull request, check that it meets these guidelines:
 3.  The pull request should work for Python 3.6 - 3.9.
     Check <https://travis-ci.org/chekos/pypums/pull_requests>
     and make sure that the tests pass for all supported Python versions.
-

docs/about/history.md

@@ -10,10 +10,10 @@
 - Bug fixes (some url's at the census website didn't have `content-size`, switched to default dict to move past that case)
 
 ## 0.0.5 (2019-05-11)
-- Add `.as_dataframe()` to ACS class. 
+- Add `.as_dataframe()` to ACS class.
 
 ## 0.0.4 (2019-05-10)
-- Add `.download_data()` to ACS class. 
+- Add `.download_data()` to ACS class.
 
 ## 0.0.3 (2019-05-09)
 - Accidentally released lol

docs/css/custom.css

@@ -3,8 +3,8 @@ div.autodoc-docstring {
     margin-bottom: 30px;
     border-left: 5px solid rgba(230, 230, 230);
   }
-  
+
   div.autodoc-members {
     padding-left: 20px;
     margin-bottom: 15px;
-  }
+  }

docs/reference/surveys.md

@@ -5,4 +5,3 @@
         members:
             - download
             - as_dataframe
-

docs/user-guide/usage.md

@@ -4,7 +4,7 @@ You can use `pypums` as a CLI on your terminal or in your scripts or notebooks.
 
 ## As a CLI
 
-On your terminal `pypums --help` or `pypums --version` to make sure it's installed 
+On your terminal `pypums --help` or `pypums --version` to make sure it's installed
 
 ```shell
 pypums --version
@@ -40,7 +40,7 @@ Options:
 ```
 
 ### `pypums download-acs`
-You can use the `download-acs` command to download the zip file containing the data for the specified ACS. It will extract it by default. 
+You can use the `download-acs` command to download the zip file containing the data for the specified ACS. It will extract it by default.
 
 This command requires at least the `year` and `state` of the survey to download.
 
@@ -98,4 +98,4 @@ ca_2023.download(data_directory="./data/")
 
 # to use it as a dataframe
 ca_2023_df = ca_2023.as_dataframe()
-```
+```

mkdocs.yml

@@ -6,7 +6,7 @@ nav:
   - User guide:
     - Installation: user-guide/installation.md
     - Usage: user-guide/usage.md
-  - API reference: 
+  - API reference:
     - Surveys (ACS): reference/surveys.md
     - CLI: reference/cli.md
   - About:
@@ -24,7 +24,7 @@ plugins:
           options:
             show_source: false
             show_root_heading: true
-  
+
 markdown_extensions:
     - smarty
     - toc:

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "us>=2.0.2",
     "httpx>=0.22.0",
     "pandas>=1.1.0",
+    "pyarrow>=10.0.0",
     "rich>=11.0.0",
     "typer>=0.4.0",
 ]

pypums/__init__.py

@@ -8,6 +8,8 @@
 from .constants import __app_name__ as __app_name__
 from .constants import __version__ as __version__
 from .decennial import get_decennial as get_decennial
+from .estimates import get_estimates as get_estimates
+from .flows import get_flows as get_flows
 from .moe import moe_product as moe_product
 from .moe import moe_prop as moe_prop
 from .moe import moe_ratio as moe_ratio

pypums/acs.py

@@ -1,9 +1,15 @@
 """American Community Survey data retrieval via the Census API."""
 
+from pathlib import Path
+
 import pandas as pd
+
 from pypums.api.client import CENSUS_API_BASE, call_census_api
 from pypums.api.geography import build_geography_query
 from pypums.api.key import census_api_key
+from pypums.cache import CensusCache
+
+_DEFAULT_CACHE_DIR = Path.home() / ".pypums" / "cache" / "api"
 
 # Z-scores for MOE confidence levels.
 _Z_SCORES: dict[int, float] = {
@@ -15,15 +21,26 @@
 # Geography columns in FIPS concatenation order.  The order matters because
 # GEOID is built by joining these columns (e.g. state+county+tract).
 _GEO_COL_ORDER = [
-    "us", "region", "division", "state", "county", "county subdivision",
-    "tract", "block group", "block", "place", "congressional district",
+    "us",
+    "region",
+    "division",
+    "state",
+    "county",
+    "county subdivision",
+    "tract",
+    "block group",
+    "block",
+    "place",
+    "congressional district",
     "state legislative district (upper chamber)",
     "state legislative district (lower chamber)",
     "zip code tabulation area",
-    "school district (unified)", "school district (elementary)",
+    "school district (unified)",
+    "school district (elementary)",
     "school district (secondary)",
     "metropolitan statistical area/micropolitan statistical area",
-    "combined statistical area", "public use microdata area",
+    "combined statistical area",
+    "public use microdata area",
     "american indian area/alaska native area/hawaiian home land",
 ]
 _GEO_COLUMNS = frozenset(_GEO_COL_ORDER)
@@ -46,6 +63,7 @@ def get_acs(
     moe_level: int = 90,
     summary_var: str | None = None,
     geometry: bool = False,
+    cache_table: bool = False,
     key: str | None = None,
 ) -> pd.DataFrame:
     """Retrieve American Community Survey data from the Census API.
@@ -74,6 +92,8 @@ def get_acs(
         Variable ID to include as denominator columns.
     geometry
         If True, return a GeoDataFrame with shapes.
+    cache_table
+        If True, cache the API response locally to avoid redundant calls.
     key
         Census API key. Falls back to ``census_api_key()``.
 
@@ -85,7 +105,9 @@ def get_acs(
     if output not in ("tidy", "wide"):
         raise ValueError(f"output must be 'tidy' or 'wide', got {output!r}")
     if moe_level not in _Z_SCORES:
-        raise ValueError(f"moe_level must be one of {sorted(_Z_SCORES)}, got {moe_level!r}")
+        raise ValueError(
+            f"moe_level must be one of {sorted(_Z_SCORES)}, got {moe_level!r}"
+        )
 
     api_key = census_api_key(key) if key else census_api_key()
     for_clause, in_clause = build_geography_query(geography, state=state, county=county)
@@ -109,6 +131,19 @@ def get_acs(
         api_vars.append(f"{summary_var}E")
         api_vars.append(f"{summary_var}M")
 
+    # Build a cache key from request parameters.
+    cache_key = (
+        f"acs_{year}_{survey}_{geography}_{state}_{county}"
+        f"_{output}_{moe_level}_{summary_var}_{','.join(api_vars)}"
+    )
+
+    # Check cache before calling API.
+    disk_cache = CensusCache(_DEFAULT_CACHE_DIR) if cache_table else None
+    if disk_cache is not None:
+        cached = disk_cache.get(cache_key)
+        if cached is not None:
+            return cached
+
     url = f"{CENSUS_API_BASE}/{year}/acs/{survey}"
     params: dict[str, str] = {
         "get": f"NAME,{','.join(api_vars)}",
@@ -179,7 +214,10 @@ def get_acs(
         # Add summary variable columns if requested.
         if summary_var is not None and summary_est_col in df.columns:
             summary_df = df[id_cols + [summary_est_col, summary_moe_col]].rename(
-                columns={summary_est_col: "summary_est", summary_moe_col: "summary_moe"},
+                columns={
+                    summary_est_col: "summary_est",
+                    summary_moe_col: "summary_moe",
+                },
             )
             result = result.merge(summary_df, on=id_cols)
 
@@ -188,4 +226,7 @@ def get_acs(
 
         result = attach_geometry(result, geography, state=state, year=year)
 
+    if disk_cache is not None:
+        disk_cache.set(cache_key, result, ttl_seconds=86400)
+
     return result

pypums/cache.py

@@ -2,19 +2,21 @@
 
 import hashlib
 import json
-import pickle
 import time
 from pathlib import Path
 
 import pandas as pd
 
-_CACHE_DATA_SUFFIX = ".pkl"
+_CACHE_DATA_SUFFIX = ".parquet"
 _CACHE_META_SUFFIX = ".meta.json"
 
 
 class CensusCache:
     """Cache Census API responses with optional TTL.
 
+    Uses Parquet for DataFrame serialization (safe to deserialize from
+    untrusted sources, unlike pickle).
+
     Parameters
     ----------
     cache_dir
@@ -53,8 +55,7 @@ def set(
         ttl_seconds
             Time-to-live in seconds. ``None`` means no expiration.
         """
-        with open(self._data_path(key), "wb") as f:
-            pickle.dump(df, f)
+        df.to_parquet(self._data_path(key))
         meta = {"created_at": time.time(), "ttl_seconds": ttl_seconds}
         self._meta_path(key).write_text(json.dumps(meta))
 
@@ -76,8 +77,7 @@ def get(self, key: str) -> pd.DataFrame | None:
                 meta_path.unlink(missing_ok=True)
                 return None
 
-        with open(data_path, "rb") as f:
-            return pickle.load(f)  # noqa: S301
+        return pd.read_parquet(data_path)
 
     def clear(self) -> None:
         """Remove all cached entries."""

pypums/datasets/__init__.py

@@ -1,3 +1,5 @@
 """Built-in reference datasets."""
 
 from pypums.datasets.fips import fips_codes as fips_codes
+from pypums.datasets.fips import lookup_fips as lookup_fips
+from pypums.datasets.fips import lookup_name as lookup_name

pypums/datasets/fips.py

@@ -7,3 +7,89 @@
 _DATA_DIR = Path(__file__).parent / "data"
 
 fips_codes: pd.DataFrame = pd.read_csv(_DATA_DIR / "fips_codes.csv", dtype=str)
+
+
+def lookup_fips(
+    state: str | None = None,
+    county: str | None = None,
+) -> str:
+    """Look up a FIPS code from a state or county name.
+
+    Parameters
+    ----------
+    state
+        State name (e.g. ``"California"``).
+    county
+        County name (e.g. ``"Los Angeles County"``). Requires *state*.
+
+    Returns
+    -------
+    str
+        The FIPS code: 2-digit state code, or state+county code.
+
+    Raises
+    ------
+    ValueError
+        If the state or county is not found.
+    """
+    if state is None:
+        raise ValueError("state is required for FIPS lookup.")
+
+    state_rows = fips_codes[fips_codes["state"].str.lower() == state.lower()]
+    if state_rows.empty:
+        raise ValueError(f"State not found: {state!r}")
+
+    state_code = state_rows["state_code"].iloc[0]
+
+    if county is None:
+        return state_code
+
+    county_rows = state_rows[state_rows["county"].str.lower() == county.lower()]
+    if county_rows.empty:
+        raise ValueError(f"County not found: {county!r} in state {state!r}")
+
+    return state_code + county_rows["county_code"].iloc[0]
+
+
+def lookup_name(
+    state_code: str | None = None,
+    county_code: str | None = None,
+) -> str:
+    """Look up a state or county name from FIPS code(s).
+
+    Parameters
+    ----------
+    state_code
+        2-digit state FIPS code (e.g. ``"06"``).
+    county_code
+        3-digit county FIPS code (e.g. ``"037"``). Requires *state_code*.
+
+    Returns
+    -------
+    str
+        The state name, or ``"County, State"`` if county_code is given.
+
+    Raises
+    ------
+    ValueError
+        If the code is not found.
+    """
+    if state_code is None:
+        raise ValueError("state_code is required for name lookup.")
+
+    state_rows = fips_codes[fips_codes["state_code"] == state_code]
+    if state_rows.empty:
+        raise ValueError(f"State code not found: {state_code!r}")
+
+    state_name = state_rows["state"].iloc[0]
+
+    if county_code is None:
+        return state_name
+
+    county_rows = state_rows[state_rows["county_code"] == county_code]
+    if county_rows.empty:
+        raise ValueError(
+            f"County code not found: {county_code!r} in state {state_code!r}"
+        )
+
+    return f"{county_rows['county'].iloc[0]}, {state_name}"