diff --git a/.github/workflows/publish-pypi.yml b/.github/workflows/publish-pypi.yml new file mode 100644 index 0000000..c5bbbd8 --- /dev/null +++ b/.github/workflows/publish-pypi.yml @@ -0,0 +1,85 @@ +# Publishes the package to PyPI when a GitHub release is published. +# +# Uses PyPI "trusted publishing" (OIDC) — no API token secrets. One-time setup +# a maintainer must do before the first release: +# 1. On pypi.org (logged in as the owning account): Account → Publishing → +# add a pending publisher with project name `spicy-regs`, owner +# `civictechdc`, repository `spicy-regs`, workflow `publish-pypi.yml`, +# environment `pypi`. +# 2. In this repo's Settings → Environments: create an environment named +# `pypi` (optionally require reviewers so releases need an approval). +# +# Release flow: bump `version` in pyproject.toml, merge to main, then create a +# GitHub release whose tag is `v` (e.g. v0.2.0). The tag/version +# guard below fails the build if the two drift. + +name: Publish to PyPI + +on: + release: + types: [published] + +jobs: + build: + runs-on: ubuntu-latest + timeout-minutes: 10 + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Check release tag matches the project version + run: | + VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml','rb'))['project']['version'])") + TAG="${GITHUB_REF_NAME}" + if [ "v${VERSION}" != "${TAG}" ]; then + echo "::error::Release tag ${TAG} does not match pyproject.toml version ${VERSION} (expected v${VERSION})." + exit 1 + fi + + - name: Build sdist and wheel + run: uv build + + - name: Smoke-test the wheel in a clean environment + run: | + uv venv /tmp/smoke + VIRTUAL_ENV=/tmp/smoke uv pip install dist/*.whl + /tmp/smoke/bin/spicy-regs --help + /tmp/smoke/bin/spicy-regs tables --source r2 + /tmp/smoke/bin/spicy-regs-dict --help + /tmp/smoke/bin/python -c "import spicy_regs.mcp_server" + + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: dist + path: dist/ + + publish: + needs: build + runs-on: ubuntu-latest + timeout-minutes: 10 + environment: + name: pypi + url: https://pypi.org/p/spicy-regs + permissions: + # Required for PyPI trusted publishing (OIDC token exchange). + id-token: write + + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + name: dist + path: dist/ + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index 06652e2..1a82cf3 100644 --- a/.gitignore +++ b/.gitignore @@ -24,4 +24,4 @@ recovery .env.local .env*.local __pycache__ -notebooks/.cache \ No newline at end of file +notebooks/.cachedist/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 56aadbb..b4f7977 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -58,12 +58,17 @@ processed output from the public R2 bucket: ```bash uv run spicy-regs download # dockets + documents + comments -uv run spicy-regs download --types comments # just comments +uv run spicy-regs download --tables comments # just comments uv run spicy-regs stats # sanity-check what you got uv run spicy-regs sample comments -n 5 # peek at a few rows +uv run spicy-regs query "SELECT count(*) FROM dockets" # arbitrary SQL (works pre-download too) ``` -Files default to `./spicy-regs-data/`; override with `-o some/dir`. +Files default to `./spicy-regs-data/`; override with `-o some/dir`. The query +commands (`tables`, `describe`, `query`, and friends) default to `--source +auto`: they read your local downloads when present and stream from the public +R2 bucket otherwise, so you can explore the full dataset before downloading +anything. **B. Run the ETL pipeline yourself (slower, but it's the real thing).** Reads JSON from the Mirrulations S3 mirror, flattens it, writes Parquet to @@ -112,7 +117,10 @@ Some terms you'll see throughout the codebase: ``` src/spicy_regs/ -├── cli.py # main CLI entrypoint (`spicy-regs` script) +├── cli/ # main CLI (`spicy-regs` script) — one module per subcommand +│ ├── _registry.py # the COMMANDS list; add your module here +│ ├── engine.py # shared DuckDB engine (views over R2 or local parquet) +│ └── download.py, query.py, … # subcommands: register() + run() ├── schemas/ # RecordType definitions — one per data shape ├── sources/ # Reader and Writer subclasses (S3, R2, parquet, …) ├── transforms/ # Transform subclasses + bulk-transform helpers @@ -175,6 +183,12 @@ not connectors — don't subclass them. ### Recipes +- **Add a CLI command:** copy `src/spicy_regs/cli/tables.py` as a template, + implement `register(subparsers)` and `run(args) -> int` in a new module under + `src/spicy_regs/cli/`, then add the module to `COMMANDS` in + `src/spicy_regs/cli/_registry.py`. Use `spicy_regs.cli.engine` for anything + that reads the published tables (it handles local vs. R2 resolution for you), + and add tests next to `tests/test_cli_commands.py`. - **Add a record shape:** construct a new `RecordType` in `schemas/` (set `path_pattern` only if your source addresses files by path). - **Add a source:** subclass `Reader`, implement `iter_records()` to yield raw @@ -197,6 +211,23 @@ not connectors — don't subclass them. - Include a test plan. - Ensure CI passes before requesting review. +## Releasing to PyPI (maintainers) + +Publishing is automated by `.github/workflows/publish-pypi.yml` via PyPI +[trusted publishing](https://docs.pypi.org/trusted-publishers/) — no API +tokens. To cut a release: + +1. Bump `version` in `pyproject.toml` and merge to `main`. +2. Create a GitHub release whose tag is `v` (e.g. `v0.2.0`). +3. The workflow builds the sdist + wheel, checks the tag matches the version, + smoke-tests the wheel in a clean environment, and publishes to PyPI. + +One-time setup (already-released projects skip this): add a pending trusted +publisher on pypi.org (project `spicy-regs`, owner `civictechdc`, repo +`spicy-regs`, workflow `publish-pypi.yml`, environment `pypi`) and create the +`pypi` environment in the repo settings. Once published, users can run the CLI +with just `uvx spicy-regs` — no git URL needed. + ## Reporting issues Open a GitHub issue with steps to reproduce, expected vs. actual behavior, and environment details. Pick a template from [the issue chooser](https://github.com/civictechdc/spicy-regs/issues/new/choose). diff --git a/README.md b/README.md index 56f88b6..d6cbb25 100644 --- a/README.md +++ b/README.md @@ -21,27 +21,50 @@ upload your output to live Cloudflare R2 storage. ### Download the published data locally -The processed dockets / documents / comments parquet files are published to a -public Cloudflare R2 bucket. Grab them with the bundled CLI — no credentials -needed: +Every published table (see the [data dictionary](https://civictechdc.github.io/spicy-regs/)) +lives as a parquet file in a public Cloudflare R2 bucket. Grab them with the +bundled CLI — no credentials needed: ```bash -uv run spicy-regs download # all three (dockets, documents, comments) -uv run spicy-regs download --types comments # comments only +uv run spicy-regs download # the core trio (dockets, documents, comments) +uv run spicy-regs download --tables comments # comments only +uv run spicy-regs download --tables feed_summary agency_stats # any published table +uv run spicy-regs download --all # everything (comments alone is multiple GB) uv run spicy-regs download -o ./my-data # custom output dir ``` -Files land in `./spicy-regs-data/` by default. Once downloaded, poke around: +Files land in `./spicy-regs-data/` by default. Downloads are atomic and +incremental — a file whose size still matches the bucket is skipped (`--force` +re-downloads). + +### Query the data with SQL + +`spicy-regs query` runs DuckDB SQL over every published table. By default +(`--source auto`) each table reads from your local download when present and +streams from the public bucket otherwise — so it works with no download at all: ```bash -uv run spicy-regs stats # row counts + top agencies per file -uv run spicy-regs sample comments -n 5 # 5 random rows from comments.parquet -uv run spicy-regs search "climate" # substring search across files +uv run spicy-regs tables # list every table + where it resolves +uv run spicy-regs describe dockets # column schema of one table +uv run spicy-regs query "SELECT agency_code, count(*) FROM dockets GROUP BY 1 ORDER BY 2 DESC LIMIT 5" +uv run spicy-regs query "SELECT * FROM feed_summary LIMIT 3" --format json # or csv +uv run spicy-regs query "SELECT * FROM agency_stats" --output stats.csv --format csv --max-rows 0 +``` + +Or poke around without writing SQL: + +```bash +uv run spicy-regs stats # row counts + top agencies per core table +uv run spicy-regs sample comments -n 5 # 5 random rows from any table +uv run spicy-regs search "climate" # substring search across the core tables uv run spicy-regs agencies # list every agency code ``` +These accept `--source local` (only your downloads) or `--source r2` (only the +bucket) when you want to pin where data comes from. + > Don't have the repo cloned? You can also run it one-shot with -> `uvx --from "spicy-regs @ git+https://github.com/civictechdc/spicy-regs" spicy-regs download --types comments`. +> `uvx --from "spicy-regs @ git+https://github.com/civictechdc/spicy-regs" spicy-regs download --tables comments`. ### Run the ETL pipeline yourself diff --git a/docs/cli.md b/docs/cli.md new file mode 100644 index 0000000..99b5401 --- /dev/null +++ b/docs/cli.md @@ -0,0 +1,112 @@ +# CLI + +The `spicy-regs` command-line tool downloads the published parquet files and +runs SQL against every table in this data dictionary — no credentials, no +setup beyond Python. + +## Install + +Run it one-shot with [uv](https://docs.astral.sh/uv/getting-started/installation/) +(nothing to install): + +```bash +uvx --from "spicy-regs @ git+https://github.com/civictechdc/spicy-regs" spicy-regs --help +``` + +Or from a clone of the repo: + +```bash +git clone https://github.com/civictechdc/spicy-regs.git +cd spicy-regs && uv sync +uv run spicy-regs --help +``` + +The examples below use `uv run spicy-regs …`; substitute the `uvx` form if you +haven't cloned the repo. + +## Where commands read data from + +Every command that reads tables takes a `--source` flag: + +| Source | Behavior | +| --- | --- | +| `auto` (default) | Per table: use your local download when present, otherwise stream from the public bucket | +| `local` | Only files downloaded to your data directory (default `./spicy-regs-data/`, override with `-o`) | +| `r2` | Only the public bucket (`https://r2.spicy-regs.dev`) | + +Because of `auto`, **you can query everything without downloading anything** — +DuckDB reads the remote parquet over HTTPS and only fetches the row groups a +query needs. Download the tables you use heavily to make repeated queries fast +and offline. + +## Explore the tables + +```bash +uv run spicy-regs tables # every table + whether it resolves locally or to R2 +uv run spicy-regs describe dockets # column names and types (matches this data dictionary) +uv run spicy-regs describe comments --format json +``` + +## Query with SQL + +`spicy-regs query` runs [DuckDB](https://duckdb.org/docs/stable/sql/introduction) +SQL with one view per published table, so you can filter, aggregate, and join +across all of them: + +```bash +# Top agencies by docket volume +uv run spicy-regs query "SELECT agency_code, count(*) AS n FROM dockets GROUP BY 1 ORDER BY n DESC LIMIT 10" + +# Join: most-commented dockets with their titles +uv run spicy-regs query " + SELECT d.docket_id, d.title, f.comment_count + FROM feed_summary f JOIN dockets d USING (docket_id) + ORDER BY f.comment_count DESC LIMIT 10" + +# Machine-readable output for scripts +uv run spicy-regs query "SELECT * FROM agency_stats LIMIT 5" --format json +uv run spicy-regs query "SELECT * FROM agency_stats" --format csv --output agency_stats.csv --max-rows 0 +``` + +Options: + +| Flag | Meaning | +| --- | --- | +| `--format table\|json\|csv` | Output format (default: aligned table) | +| `--max-rows N` | Cap returned rows; `0` = unlimited (default 25) | +| `--output FILE` | Write results to a file instead of stdout | +| `--source`, `--r2-url`, `-o` | Data source controls (see above) | + +Keep a `LIMIT` on exploratory queries — `comments` in particular is tens of +millions of rows. + +## Quick looks without SQL + +```bash +uv run spicy-regs stats # row counts + top agencies for the core tables +uv run spicy-regs sample comments -n 5 # random rows from any table (--agency EPA to filter) +uv run spicy-regs search "climate" # substring search across dockets/documents/comments +uv run spicy-regs agencies # every agency code in the dataset +``` + +## Download the parquet files + +```bash +uv run spicy-regs download # core trio: dockets, documents, comments +uv run spicy-regs download --tables feed_summary agency_stats +uv run spicy-regs download --all # every table (comments alone is multiple GB) +uv run spicy-regs download -o ./my-data # custom directory +``` + +Downloads stream with a progress bar and are written atomically, so an +interrupted download never leaves a truncated file. Files whose size still +matches the bucket are skipped on re-runs; `--force` re-downloads +unconditionally. + +## Extending the CLI + +Each subcommand is a small module in +[`src/spicy_regs/cli/`](https://github.com/civictechdc/spicy-regs/tree/main/src/spicy_regs/cli) +— adding one is a copy-a-template-and-register-it change. See the +[contributing guide](https://github.com/civictechdc/spicy-regs/blob/main/CONTRIBUTING.md) +for the recipe. diff --git a/docs/index.md b/docs/index.md index 92020a0..1a6bf7d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -127,11 +127,19 @@ of `dockets`; they join to the corpus (and to each other) on a few shared keys: === "CLI" + The bundled [`spicy-regs` CLI](cli.md) downloads the parquet files and runs + SQL over every table — by default it reads local downloads when present and + streams from the bucket otherwise, so no download is required: + ```bash - uvx --from "spicy-regs @ git+https://github.com/civictechdc/spicy-regs" spicy-regs download - uv run spicy-regs stats + uvx --from "spicy-regs @ git+https://github.com/civictechdc/spicy-regs" \ + spicy-regs query "SELECT agency_code, count(*) FROM dockets GROUP BY 1 ORDER BY 2 DESC LIMIT 5" + uv run spicy-regs tables # list every table + uv run spicy-regs download --all # take the dataset offline ``` + See the [CLI page](cli.md) for the full command reference. + === "DuckDB (SQL)" ```sql diff --git a/mkdocs.yml b/mkdocs.yml index 96117e3..91ef59d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -36,6 +36,7 @@ markdown_extensions: nav: - Overview: index.md + - CLI: cli.md - Querying with Python: querying-python.md - Tables: - dockets: tables/dockets.md diff --git a/pyproject.toml b/pyproject.toml index ce69621..3e8c3da 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,6 +2,17 @@ name = "spicy-regs" version = "0.1.0" description = "Regulations Data Analysis Tools" +readme = "README.md" +license = "MIT" +license-files = ["LICENSE"] +authors = [{ name = "Civic Tech DC" }] +keywords = ["regulations", "rulemaking", "regulations.gov", "civic-tech", "open-data", "duckdb", "parquet"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering :: Information Analysis", +] requires-python = ">=3.10" dependencies = [ "boto3", @@ -42,6 +53,12 @@ docs = [ "mkdocs-material", ] +[project.urls] +Homepage = "https://github.com/civictechdc/spicy-regs" +Documentation = "https://docs.spicy-regs.dev" +Repository = "https://github.com/civictechdc/spicy-regs" +Issues = "https://github.com/civictechdc/spicy-regs/issues" + [project.scripts] run-pipeline = "spicy_regs.pipelines.regulations:app" run-rollup-feed-summary = "spicy_regs.pipelines.rollups.feed_summary:app" diff --git a/src/spicy_regs/cli.py b/src/spicy_regs/cli.py deleted file mode 100644 index e7e1226..0000000 --- a/src/spicy_regs/cli.py +++ /dev/null @@ -1,256 +0,0 @@ -#!/usr/bin/env python3 -""" -Spicy Regs CLI - Download and explore federal regulations data. - -Usage: - uvx spicy-regs download # Download all parquet files - uvx spicy-regs stats # Show dataset statistics - uvx spicy-regs sample dockets # Show sample rows from a dataset - uvx spicy-regs search "climate" # Search across datasets -""" - -import argparse -import sys -from pathlib import Path -from urllib.request import urlretrieve -from urllib.error import URLError - -# Public URL for the R2 bucket -PUBLIC_URL = "https://r2.spicy-regs.dev" -DATA_TYPES = ["dockets", "documents", "comments", "manifest"] -DEFAULT_OUTPUT_DIR = Path("./spicy-regs-data") - - -def get_output_dir(args) -> Path: - """Get output directory from args or default.""" - output_dir = Path(args.output_dir) if hasattr(args, "output_dir") and args.output_dir else DEFAULT_OUTPUT_DIR - output_dir.mkdir(parents=True, exist_ok=True) - return output_dir - - -def download_file(name: str, output_dir: Path, force: bool = False) -> Path | None: - """Download a parquet file from R2.""" - url = f"{PUBLIC_URL}/{name}.parquet" - local_path = output_dir / f"{name}.parquet" - - if local_path.exists() and not force: - size_mb = local_path.stat().st_size / (1024 * 1024) - print(f" ✓ {name}.parquet already exists ({size_mb:.1f} MB)") - return local_path - - print(f" ⬇ Downloading {name}.parquet...") - try: - urlretrieve(url, local_path) - size_mb = local_path.stat().st_size / (1024 * 1024) - print(f" ✓ {name}.parquet ({size_mb:.1f} MB)") - return local_path - except URLError as e: - print(f" ✗ Failed to download {name}.parquet: {e}") - return None - - -def cmd_download(args): - """Download parquet files from R2.""" - output_dir = get_output_dir(args) - print(f"Downloading to: {output_dir.absolute()}") - - types_to_download = args.types if args.types else ["dockets", "documents", "comments"] - - for data_type in types_to_download: - download_file(data_type, output_dir, force=args.force) - - print(f"\nDone! Data saved to: {output_dir.absolute()}") - - -def cmd_stats(args): - """Show statistics for downloaded datasets.""" - try: - import polars as pl - except ImportError: - print("Please install polars: pip install polars") - sys.exit(1) - - output_dir = get_output_dir(args) - - print("=" * 60) - print("Dataset Statistics") - print("=" * 60) - - for data_type in ["dockets", "documents", "comments"]: - parquet_file = output_dir / f"{data_type}.parquet" - if not parquet_file.exists(): - print(f"\n{data_type.upper()}: Not downloaded yet (run: spicy-regs download)") - continue - - df = pl.read_parquet(parquet_file) - size_mb = parquet_file.stat().st_size / (1024 * 1024) - - print(f"\n{data_type.upper()} ({size_mb:.1f} MB)") - print("-" * 40) - print(f" Rows: {len(df):,}") - print(f" Columns: {', '.join(df.columns)}") - - # Agency breakdown - if "agency_code" in df.columns: - agency_counts = df.group_by("agency_code").len().sort("len", descending=True) - top_agencies = agency_counts.head(5) - print(" Top agencies:") - for row in top_agencies.iter_rows(): - print(f" {row[0]}: {row[1]:,}") - - -def cmd_sample(args): - """Show sample rows from a dataset.""" - try: - import polars as pl - except ImportError: - print("Please install polars: pip install polars") - sys.exit(1) - - output_dir = get_output_dir(args) - parquet_file = output_dir / f"{args.data_type}.parquet" - - if not parquet_file.exists(): - print(f"File not found: {parquet_file}") - print("Run: spicy-regs download") - sys.exit(1) - - df = pl.read_parquet(parquet_file) - - if args.agency: - df = df.filter(pl.col("agency_code") == args.agency) - - sample = df.sample(min(args.n, len(df))) - - print(f"\nSample from {args.data_type} ({len(df):,} total rows):") - print("=" * 80) - print(sample) - - -def cmd_search(args): - """Search across datasets.""" - try: - import polars as pl - except ImportError: - print("Please install polars: pip install polars") - sys.exit(1) - - output_dir = get_output_dir(args) - query = args.query.lower() - - print(f"Searching for: '{args.query}'") - print("=" * 60) - - search_configs = { - "dockets": ["title", "abstract"], - "documents": ["title", "text_content"], - "comments": ["title", "comment", "text_content"], - } - - for data_type, columns in search_configs.items(): - parquet_file = output_dir / f"{data_type}.parquet" - if not parquet_file.exists(): - continue - - df = pl.read_parquet(parquet_file) - - # Build filter for any column containing the query - filters = None - for col in columns: - if col in df.columns: - col_filter = pl.col(col).str.to_lowercase().str.contains(query, literal=True) - filters = col_filter if filters is None else (filters | col_filter) - - if filters is not None: - matches = df.filter(filters) - if len(matches) > 0: - print(f"\n{data_type.upper()}: {len(matches):,} matches") - print("-" * 40) - sample = matches.head(args.limit) - for row in sample.iter_rows(named=True): - id_col = list(row.keys())[0] - title = row.get("title", "")[:80] if row.get("title") else "(no title)" - print(f" {row[id_col]}: {title}") - - -def cmd_agencies(args): - """List all agencies in the dataset.""" - try: - import polars as pl - except ImportError: - print("Please install polars: pip install polars") - sys.exit(1) - - output_dir = get_output_dir(args) - - # Try to get agency list from any available file - for data_type in ["dockets", "documents", "comments"]: - parquet_file = output_dir / f"{data_type}.parquet" - if parquet_file.exists(): - df = pl.read_parquet(parquet_file, columns=["agency_code"]) - agencies = df["agency_code"].unique().sort().to_list() - - print(f"Agencies ({len(agencies)} total):") - print("=" * 40) - for agency in agencies: - if agency: - print(f" {agency}") - return - - print("No data downloaded yet. Run: spicy-regs download") - - -def main(): - parser = argparse.ArgumentParser( - prog="spicy-regs", - description="Download and explore federal regulations data from Spicy Regs", - ) - parser.add_argument( - "--output-dir", - "-o", - help=f"Output directory for data files (default: {DEFAULT_OUTPUT_DIR})", - default=None, - ) - - subparsers = parser.add_subparsers(dest="command", help="Available commands") - - # Download command - download_parser = subparsers.add_parser("download", help="Download parquet files") - download_parser.add_argument("--force", "-f", action="store_true", help="Force re-download") - download_parser.add_argument( - "--types", nargs="+", choices=["dockets", "documents", "comments"], help="Specific data types to download" - ) - download_parser.set_defaults(func=cmd_download) - - # Stats command - stats_parser = subparsers.add_parser("stats", help="Show dataset statistics") - stats_parser.set_defaults(func=cmd_stats) - - # Sample command - sample_parser = subparsers.add_parser("sample", help="Show sample rows") - sample_parser.add_argument("data_type", choices=["dockets", "documents", "comments"]) - sample_parser.add_argument("-n", type=int, default=5, help="Number of rows") - sample_parser.add_argument("--agency", help="Filter by agency code") - sample_parser.set_defaults(func=cmd_sample) - - # Search command - search_parser = subparsers.add_parser("search", help="Search across datasets") - search_parser.add_argument("query", help="Search query") - search_parser.add_argument("--limit", "-l", type=int, default=10, help="Max results per type") - search_parser.set_defaults(func=cmd_search) - - # Agencies command - agencies_parser = subparsers.add_parser("agencies", help="List all agencies") - agencies_parser.set_defaults(func=cmd_agencies) - - args = parser.parse_args() - - if args.command is None: - parser.print_help() - sys.exit(0) - - args.func(args) - - -if __name__ == "__main__": - main() diff --git a/src/spicy_regs/cli/__init__.py b/src/spicy_regs/cli/__init__.py new file mode 100644 index 0000000..9a5076f --- /dev/null +++ b/src/spicy_regs/cli/__init__.py @@ -0,0 +1,52 @@ +"""Spicy Regs CLI — download and query federal regulations data. + +Usage: + spicy-regs download # fetch the core parquet files + spicy-regs tables # list every queryable table + spicy-regs describe dockets # column schema of one table + spicy-regs query "SELECT ..." # run SQL (local files or straight off R2) + +Each subcommand lives in its own module in this package; see ``_registry.py`` +for how to add one. +""" + +from __future__ import annotations + +import argparse +import sys + +from spicy_regs.cli._common import DEFAULT_OUTPUT_DIR + + +def build_parser() -> argparse.ArgumentParser: + # Imported here (not at module top) so command modules can import helpers + # from ``spicy_regs.cli._common`` without a circular import at package load. + from spicy_regs.cli._registry import COMMANDS + + parser = argparse.ArgumentParser( + prog="spicy-regs", + description="Download and explore federal regulations data from Spicy Regs", + ) + parser.add_argument( + "--output-dir", + "-o", + default=None, + help=f"Directory for local data files (default: {DEFAULT_OUTPUT_DIR})", + ) + subparsers = parser.add_subparsers(dest="command", help="Available commands") + for module in COMMANDS: + module.register(subparsers) + return parser + + +def main(argv: list[str] | None = None) -> None: + parser = build_parser() + args = parser.parse_args(argv) + if args.command is None: + parser.print_help() + sys.exit(0) + sys.exit(args.run(args)) + + +if __name__ == "__main__": + main() diff --git a/src/spicy_regs/cli/_common.py b/src/spicy_regs/cli/_common.py new file mode 100644 index 0000000..265fe26 --- /dev/null +++ b/src/spicy_regs/cli/_common.py @@ -0,0 +1,47 @@ +"""Helpers shared by the ``spicy-regs`` command modules.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from spicy_regs.data_dictionary import DEFAULT_R2_BASE_URL + +DEFAULT_OUTPUT_DIR = Path("./spicy-regs-data") + + +def get_output_dir(args: argparse.Namespace) -> Path: + """Resolve the data directory from ``--output-dir`` (global or per-command).""" + raw = getattr(args, "output_dir", None) + return Path(raw) if raw else DEFAULT_OUTPUT_DIR + + +def add_output_dir_argument(parser: argparse.ArgumentParser) -> None: + """Accept ``-o/--output-dir`` after the subcommand too (``spicy-regs download -o dir``). + + ``default=SUPPRESS`` keeps the subparser from clobbering a value given + before the subcommand (``spicy-regs -o dir download``) — argparse subparser + defaults would otherwise overwrite the already-parsed global value. + """ + parser.add_argument( + "--output-dir", + "-o", + default=argparse.SUPPRESS, + help=f"Directory for local data files (default: {DEFAULT_OUTPUT_DIR})", + ) + + +def add_source_arguments(parser: argparse.ArgumentParser) -> None: + """Standard flags for commands that read tables through the query engine.""" + parser.add_argument( + "--source", + choices=["r2", "local", "auto"], + default="auto", + help="Read from the public R2 bucket, local downloads, or per-table whichever is present locally (default)", + ) + parser.add_argument( + "--r2-url", + default=DEFAULT_R2_BASE_URL, + help=f"Base URL of the public parquet bucket (default: {DEFAULT_R2_BASE_URL})", + ) + add_output_dir_argument(parser) diff --git a/src/spicy_regs/cli/_output.py b/src/spicy_regs/cli/_output.py new file mode 100644 index 0000000..05a793d --- /dev/null +++ b/src/spicy_regs/cli/_output.py @@ -0,0 +1,98 @@ +"""Result formatting for the ``spicy-regs`` CLI: table, JSON, and CSV.""" + +from __future__ import annotations + +import csv +import json +from datetime import date, datetime, time, timedelta +from decimal import Decimal +from typing import IO, TYPE_CHECKING, Any +from uuid import UUID + +if TYPE_CHECKING: + from spicy_regs.cli.engine import QueryResult + +# Cap cell width in table output so a full-text comment column doesn't wrap the +# whole terminal; --format json/csv always carries complete values. +MAX_CELL_WIDTH = 80 + +FORMATS = ("table", "json", "csv") + + +def jsonify(value: Any) -> Any: + """Coerce DuckDB row values into JSON-serializable forms. + + Same coercions as ``spicy_regs.mcp_server._jsonify``; duplicated because the + MCP server must stay self-contained for its Vercel sync copy. + """ + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, (datetime, date, time)): + return value.isoformat() + if isinstance(value, timedelta): + return value.total_seconds() + if isinstance(value, Decimal): + return str(value) + if isinstance(value, UUID): + return str(value) + if isinstance(value, (bytes, bytearray, memoryview)): + return bytes(value).hex() + if isinstance(value, (list, tuple, set, frozenset)): + return [jsonify(v) for v in value] + if isinstance(value, dict): + return {str(k): jsonify(v) for k, v in value.items()} + return str(value) + + +def _cell(value: Any) -> str: + if value is None: + return "" + text = str(jsonify(value)).replace("\n", " ").replace("\r", " ") + if len(text) > MAX_CELL_WIDTH: + return text[: MAX_CELL_WIDTH - 1] + "…" + return text + + +def format_table(result: QueryResult) -> str: + """Render rows as a width-aligned text table.""" + if not result.columns: + return "(no columns)" + cells = [[_cell(v) for v in row] for row in result.rows] + widths = [len(c) for c in result.columns] + for row in cells: + for i, text in enumerate(row): + widths[i] = max(widths[i], len(text)) + header = " | ".join(name.ljust(widths[i]) for i, name in enumerate(result.columns)) + rule = "-+-".join("-" * w for w in widths) + lines = [header, rule] + lines.extend(" | ".join(text.ljust(widths[i]) for i, text in enumerate(row)) for row in cells) + if result.truncated: + lines.append(f"({len(result.rows)} rows shown; more available — raise --max-rows or add a LIMIT)") + else: + lines.append(f"({len(result.rows)} row{'s' if len(result.rows) != 1 else ''})") + return "\n".join(lines) + + +def write_json(result: QueryResult, fh: IO[str]) -> None: + """Write rows as a JSON array of {column: value} objects.""" + rows = [{col: jsonify(val) for col, val in zip(result.columns, row)} for row in result.rows] + json.dump(rows, fh, indent=2) + fh.write("\n") + + +def write_csv(result: QueryResult, fh: IO[str]) -> None: + """Write rows as CSV with a header line.""" + writer = csv.writer(fh) + writer.writerow(result.columns) + for row in result.rows: + writer.writerow(["" if v is None else jsonify(v) for v in row]) + + +def write_result(result: QueryResult, fmt: str, fh: IO[str]) -> None: + """Dispatch to the writer for ``fmt`` (one of :data:`FORMATS`).""" + if fmt == "json": + write_json(result, fh) + elif fmt == "csv": + write_csv(result, fh) + else: + fh.write(format_table(result) + "\n") diff --git a/src/spicy_regs/cli/_registry.py b/src/spicy_regs/cli/_registry.py new file mode 100644 index 0000000..326a911 --- /dev/null +++ b/src/spicy_regs/cli/_registry.py @@ -0,0 +1,12 @@ +"""The ``spicy-regs`` command registry. + +To add a command: create a module in this package with ``register(subparsers)`` +and ``run(args) -> int`` functions (copy ``tables.py`` as a template), then add +the module here. The list order is the help-text order. +""" + +from __future__ import annotations + +from spicy_regs.cli import agencies, describe, download, query, sample, search, stats, tables + +COMMANDS = [download, tables, describe, query, stats, sample, search, agencies] diff --git a/src/spicy_regs/cli/agencies.py b/src/spicy_regs/cli/agencies.py new file mode 100644 index 0000000..dcf7734 --- /dev/null +++ b/src/spicy_regs/cli/agencies.py @@ -0,0 +1,42 @@ +"""``spicy-regs agencies`` — list every agency code in the dataset.""" + +from __future__ import annotations + +import argparse +import sys + +import duckdb + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir + +# Checked in order; dockets is the smallest table that carries every agency. +CANDIDATE_TABLES = ("dockets", "documents", "comments") + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("agencies", help="List all agency codes") + add_source_arguments(parser) + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + for table in CANDIDATE_TABLES: + spec = specs.get(table) + if spec is None: + continue + con = engine.connect({table: spec}) + sql = f"SELECT DISTINCT agency_code FROM {table} WHERE agency_code IS NOT NULL ORDER BY 1" + try: + result = engine.run_query(con, sql, max_rows=0) + except duckdb.Error as exc: + print(f"Could not read {table} at {spec.location}: {exc}", file=sys.stderr) + continue + print(f"Agencies ({len(result.rows)} total, from {table}):") + print("=" * 40) + for (agency,) in result.rows: + print(f" {agency}") + return 0 + print("No data available. Run: spicy-regs download", file=sys.stderr) + return 1 diff --git a/src/spicy_regs/cli/describe.py b/src/spicy_regs/cli/describe.py new file mode 100644 index 0000000..c144777 --- /dev/null +++ b/src/spicy_regs/cli/describe.py @@ -0,0 +1,38 @@ +"""``spicy-regs describe`` — show the column schema of one table.""" + +from __future__ import annotations + +import argparse +import sys + +import duckdb + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir +from spicy_regs.cli._output import write_result + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("describe", help="Show the column schema of a table") + parser.add_argument("table", choices=engine.TABLES, metavar="TABLE", help=f"One of: {', '.join(engine.TABLES)}") + add_source_arguments(parser) + parser.add_argument("--format", choices=["table", "json"], default="table", help="Output format") + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + spec = specs.get(args.table) + if spec is None: + print(f"Table '{args.table}' is not available locally. Run: spicy-regs download", file=sys.stderr) + return 1 + # Only bind the one view we need — binding reads parquet metadata, which is + # an HTTP round-trip per table for remote sources. + con = engine.connect({args.table: spec}) + try: + result = engine.run_query(con, f"DESCRIBE {args.table}", max_rows=0) + except duckdb.Error as exc: + print(f"Could not describe {args.table} at {spec.location}: {exc}", file=sys.stderr) + return 1 + write_result(result, args.format, sys.stdout) + return 0 diff --git a/src/spicy_regs/cli/download.py b/src/spicy_regs/cli/download.py new file mode 100644 index 0000000..386fc23 --- /dev/null +++ b/src/spicy_regs/cli/download.py @@ -0,0 +1,127 @@ +"""``spicy-regs download`` — fetch published parquet files from the public bucket. + +Downloads are streamed to a ``.tmp`` file and atomically renamed into place, so +an interrupted download never leaves a truncated parquet behind. Files that +already match the remote size (HEAD Content-Length) are skipped; ``--force`` +re-downloads unconditionally. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import httpx +from tqdm import tqdm + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_output_dir_argument, get_output_dir +from spicy_regs.data_dictionary import DEFAULT_R2_BASE_URL + +# Everything published to the bucket: the queryable tables plus the pipeline +# manifest snapshot. +DOWNLOADABLE = (*engine.TABLES, "manifest") +DEFAULT_TABLES = ("dockets", "documents", "comments") + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("download", help="Download parquet files from the public bucket") + parser.add_argument( + "--tables", + nargs="+", + choices=DOWNLOADABLE, + metavar="TABLE", + help=f"Tables to download (default: {' '.join(DEFAULT_TABLES)}); one of: {', '.join(DOWNLOADABLE)}", + ) + # Backwards-compatible alias for the pre-1.0 flag spelling. + parser.add_argument("--types", nargs="+", choices=DOWNLOADABLE, dest="tables", help=argparse.SUPPRESS) + parser.add_argument( + "--all", + action="store_true", + dest="download_all", + help="Download every published table (comments alone is multiple GB)", + ) + parser.add_argument("--force", "-f", action="store_true", help="Re-download even if the local file looks current") + parser.add_argument( + "--r2-url", + default=DEFAULT_R2_BASE_URL, + help=f"Base URL of the public parquet bucket (default: {DEFAULT_R2_BASE_URL})", + ) + add_output_dir_argument(parser) + parser.set_defaults(run=run) + + +def _build_client(transport: httpx.BaseTransport | None = None) -> httpx.Client: + """HTTP client for bucket downloads (``transport`` is a test seam).""" + return httpx.Client(follow_redirects=True, timeout=httpx.Timeout(30.0, read=300.0), transport=transport) + + +def _is_up_to_date(client: httpx.Client, url: str, local_path: Path) -> bool: + """True when the local file's size matches the remote Content-Length.""" + if not local_path.exists(): + return False + try: + response = client.head(url) + response.raise_for_status() + content_length = response.headers.get("content-length") + except httpx.HTTPError: + return False # can't tell — re-download + return content_length is not None and int(content_length) == local_path.stat().st_size + + +def download_file(client: httpx.Client, url: str, local_path: Path, force: bool = False) -> str: + """Download one file; returns ``"downloaded"``, ``"skipped"``, ``"missing"``, or ``"failed"``.""" + name = local_path.name + if not force and _is_up_to_date(client, url, local_path): + size_mb = local_path.stat().st_size / (1024 * 1024) + print(f" ✓ {name} up to date ({size_mb:.1f} MB)") + return "skipped" + + temp_path = local_path.with_suffix(local_path.suffix + ".tmp") + try: + with client.stream("GET", url) as response: + if response.status_code == 404: + print(f" - {name} is not published at {url}; skipping") + return "missing" + response.raise_for_status() + total = int(response.headers.get("content-length", 0)) or None + with ( + temp_path.open("wb") as fh, + tqdm(total=total, unit="B", unit_scale=True, desc=f" ⬇ {name}", disable=None, leave=False) as bar, + ): + for chunk in response.iter_bytes(): + fh.write(chunk) + bar.update(len(chunk)) + temp_path.replace(local_path) + except (httpx.HTTPError, OSError) as exc: + temp_path.unlink(missing_ok=True) + print(f" ✗ Failed to download {name}: {exc}", file=sys.stderr) + return "failed" + size_mb = local_path.stat().st_size / (1024 * 1024) + print(f" ✓ {name} ({size_mb:.1f} MB)") + return "downloaded" + + +def run(args: argparse.Namespace) -> int: + output_dir = get_output_dir(args) + output_dir.mkdir(parents=True, exist_ok=True) + print(f"Downloading to: {output_dir.absolute()}") + + if args.download_all: + tables = DOWNLOADABLE + elif args.tables: + tables = tuple(dict.fromkeys(args.tables)) # de-dupe, keep order + else: + tables = DEFAULT_TABLES + + base_url = args.r2_url.rstrip("/") + failures = 0 + with _build_client() as client: + for table in tables: + status = download_file(client, f"{base_url}/{table}.parquet", output_dir / f"{table}.parquet", args.force) + if status == "failed": + failures += 1 + + print(f"\nDone! Data saved to: {output_dir.absolute()}") + return 1 if failures else 0 diff --git a/src/spicy_regs/cli/engine.py b/src/spicy_regs/cli/engine.py new file mode 100644 index 0000000..13829e6 --- /dev/null +++ b/src/spicy_regs/cli/engine.py @@ -0,0 +1,144 @@ +"""DuckDB query engine shared by the ``spicy-regs`` CLI commands. + +Builds one DuckDB view per published table, pointing either at the public R2 +parquet (``https://r2.spicy-regs.dev/.parquet``) or at files downloaded +by ``spicy-regs download``. The connect/skip patterns mirror +``spicy_regs.mcp_server`` (which must stay self-contained for its Vercel sync +copy, so the small overlap is deliberate); the table list itself is imported +from ``spicy_regs.data_dictionary`` so it has exactly one owner. +""" + +from __future__ import annotations + +import sys +import tempfile +from dataclasses import dataclass +from pathlib import Path + +import duckdb + +from spicy_regs.data_dictionary import DEFAULT_R2_BASE_URL, TABLES + +__all__ = [ + "DEFAULT_R2_BASE_URL", + "TABLES", + "QueryResult", + "ViewSpec", + "connect", + "escape_sql_string", + "local_view_specs", + "remote_view_specs", + "resolve_view_specs", + "run_query", +] + + +@dataclass(frozen=True) +class ViewSpec: + """Where one logical table's view reads from.""" + + table: str + kind: str # "local" | "r2" + location: str # file path or URL + sql: str # SELECT statement the view is created from + + +@dataclass +class QueryResult: + columns: list[str] + rows: list[tuple] + truncated: bool # True when rows were cut off at max_rows + + +def escape_sql_string(value: str) -> str: + """Escape a value for inlining into a single-quoted SQL string literal.""" + return value.replace("'", "''") + + +def _read_parquet_sql(location: str, **options: bool) -> str: + opts = "".join(f", {key}={str(val).lower()}" for key, val in options.items()) + return f"SELECT * FROM read_parquet('{escape_sql_string(location)}'{opts})" + + +def local_view_specs(data_dir: Path) -> dict[str, ViewSpec]: + """Specs for every published table that exists under ``data_dir``. + + Besides the flat ``
.parquet`` files that ``spicy-regs download`` + writes, ``comments`` may exist as the pipeline's partitioned + ``comments/**/*.parquet`` tree — handled the same way as the plugin's + standalone query script. + """ + specs: dict[str, ViewSpec] = {} + for table in TABLES: + flat = data_dir / f"{table}.parquet" + if flat.exists(): + specs[table] = ViewSpec(table, "local", str(flat), _read_parquet_sql(str(flat))) + continue + if table == "comments": + partitioned = data_dir / "comments" + if partitioned.is_dir(): + pattern = str(partitioned / "**" / "*.parquet") + sql = _read_parquet_sql(pattern, union_by_name=True, hive_partitioning=True) + specs[table] = ViewSpec(table, "local", pattern, sql) + return specs + + +def remote_view_specs(base_url: str = DEFAULT_R2_BASE_URL) -> dict[str, ViewSpec]: + """Specs for all published tables against the public bucket.""" + url = base_url.rstrip("/") + return { + table: ViewSpec(table, "r2", f"{url}/{table}.parquet", _read_parquet_sql(f"{url}/{table}.parquet")) + for table in TABLES + } + + +def resolve_view_specs(source: str, data_dir: Path, base_url: str = DEFAULT_R2_BASE_URL) -> dict[str, ViewSpec]: + """Resolve where each table should be read from. + + ``source`` is ``"r2"``, ``"local"``, or ``"auto"`` — auto prefers the local + copy of each table when present and falls back to R2 for the rest, so a + partial download still gives full query coverage. + """ + if source == "local": + return local_view_specs(data_dir) + if source == "r2": + return remote_view_specs(base_url) + if source == "auto": + specs = remote_view_specs(base_url) + specs.update(local_view_specs(data_dir)) + return specs + raise ValueError(f"Unknown source {source!r}; expected 'r2', 'local', or 'auto'") + + +def connect(specs: dict[str, ViewSpec]) -> duckdb.DuckDBPyConnection: + """Open an in-memory DuckDB connection with one view per spec. + + A view whose parquet can't be read (table registered but not published yet, + local file corrupt, network error) is skipped with a warning instead of + failing the whole command — same degradation as the MCP server. + """ + con = duckdb.connect() + # Must precede INSTALL/LOAD: the extension cache lands under + # /.duckdb, and $HOME may be unset or read-only. + con.execute(f"SET home_directory='{escape_sql_string(tempfile.gettempdir())}'") + con.execute("SET preserve_insertion_order=false") + if any(spec.kind == "r2" for spec in specs.values()): + con.execute("INSTALL httpfs") + con.execute("LOAD httpfs") + for spec in specs.values(): + try: + con.execute(f"CREATE VIEW {spec.table} AS {spec.sql}") + except duckdb.Error as exc: + print(f"warning: table {spec.table} not available at {spec.location}: {exc}", file=sys.stderr) + return con + + +def run_query(con: duckdb.DuckDBPyConnection, sql: str, max_rows: int) -> QueryResult: + """Execute ``sql`` and return up to ``max_rows`` rows (``0`` = unlimited).""" + cursor = con.execute(sql) + columns = [desc[0] for desc in cursor.description] if cursor.description else [] + if max_rows <= 0: + return QueryResult(columns=columns, rows=cursor.fetchall(), truncated=False) + rows = cursor.fetchmany(max_rows) + truncated = len(rows) == max_rows and cursor.fetchone() is not None + return QueryResult(columns=columns, rows=rows, truncated=truncated) diff --git a/src/spicy_regs/cli/query.py b/src/spicy_regs/cli/query.py new file mode 100644 index 0000000..9c4c12d --- /dev/null +++ b/src/spicy_regs/cli/query.py @@ -0,0 +1,56 @@ +"""``spicy-regs query`` — run SQL against the published tables. + +Every published table is available as a view (see ``spicy-regs tables``), so +queries can join across them, e.g.:: + + spicy-regs query "SELECT agency_code, count(*) FROM dockets GROUP BY 1 ORDER BY 2 DESC LIMIT 5" +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import duckdb + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir +from spicy_regs.cli._output import FORMATS, write_result + +DEFAULT_MAX_ROWS = 25 + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("query", help="Run a SQL query against the tables") + parser.add_argument("sql", help='SQL to run, e.g. "SELECT * FROM dockets LIMIT 5"') + add_source_arguments(parser) + parser.add_argument("--format", choices=FORMATS, default="table", help="Output format (default: table)") + parser.add_argument( + "--max-rows", + type=int, + default=DEFAULT_MAX_ROWS, + help=f"Maximum rows to return; 0 = unlimited (default: {DEFAULT_MAX_ROWS})", + ) + parser.add_argument("--output", metavar="FILE", default=None, help="Write results to FILE instead of stdout") + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + if not specs: + print("No tables available for this source. Run: spicy-regs download", file=sys.stderr) + return 1 + con = engine.connect(specs) + try: + result = engine.run_query(con, args.sql, args.max_rows) + except duckdb.Error as exc: + print(f"Query failed: {exc}", file=sys.stderr) + return 1 + if args.output: + with Path(args.output).open("w", newline="") as fh: + write_result(result, args.format, fh) + print(f"Wrote {len(result.rows)} rows to {args.output}") + else: + write_result(result, args.format, sys.stdout) + return 0 diff --git a/src/spicy_regs/cli/sample.py b/src/spicy_regs/cli/sample.py new file mode 100644 index 0000000..1761364 --- /dev/null +++ b/src/spicy_regs/cli/sample.py @@ -0,0 +1,41 @@ +"""``spicy-regs sample`` — show random rows from a table.""" + +from __future__ import annotations + +import argparse +import sys + +import duckdb + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir +from spicy_regs.cli._output import format_table + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("sample", help="Show random sample rows from a table") + parser.add_argument("data_type", choices=engine.TABLES, metavar="TABLE", help=f"One of: {', '.join(engine.TABLES)}") + parser.add_argument("-n", type=int, default=5, help="Number of rows (default: 5)") + parser.add_argument("--agency", help="Filter by agency code") + add_source_arguments(parser) + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + table = args.data_type + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + spec = specs.get(table) + if spec is None: + print(f"Table '{table}' is not available locally. Run: spicy-regs download", file=sys.stderr) + return 1 + con = engine.connect({table: spec}) + where = f"WHERE agency_code = '{engine.escape_sql_string(args.agency)}'" if args.agency else "" + try: + result = engine.run_query(con, f"SELECT * FROM {table} {where} USING SAMPLE {int(args.n)} ROWS", max_rows=0) + except duckdb.Error as exc: + print(f"Could not sample {table}: {exc}", file=sys.stderr) + return 1 + print(f"\nSample from {table} ({spec.kind}: {spec.location}):") + print("=" * 80) + print(format_table(result)) + return 0 diff --git a/src/spicy_regs/cli/search.py b/src/spicy_regs/cli/search.py new file mode 100644 index 0000000..bcb6663 --- /dev/null +++ b/src/spicy_regs/cli/search.py @@ -0,0 +1,71 @@ +"""``spicy-regs search`` — substring search across the core tables.""" + +from __future__ import annotations + +import argparse +import sys + +import duckdb + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir + +# Table -> (id column, text columns searched). Add an entry here to make +# another table searchable. +SEARCH_CONFIGS: dict[str, tuple[str, tuple[str, ...]]] = { + "dockets": ("docket_id", ("title", "abstract")), + "documents": ("document_id", ("title", "text_content")), + "comments": ("comment_id", ("title", "comment", "text_content")), +} + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("search", help="Search for a substring across the core tables") + parser.add_argument("query", help="Text to search for (case-insensitive substring)") + parser.add_argument("--limit", "-l", type=int, default=10, help="Max results per table (default: 10)") + add_source_arguments(parser) + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + needle = engine.escape_sql_string(args.query.lower()) + + print(f"Searching for: '{args.query}'") + print("=" * 60) + + failures = 0 + for table, (id_column, text_columns) in SEARCH_CONFIGS.items(): + spec = specs.get(table) + if spec is None: + continue + con = engine.connect({table: spec}) + try: + available = {row[0] for row in engine.run_query(con, f"DESCRIBE {table}", max_rows=0).rows} + except duckdb.Error as exc: + print(f"\n{table.upper()}: could not read {spec.location}: {exc}", file=sys.stderr) + failures += 1 + continue + columns = [c for c in text_columns if c in available] + if not columns: + continue + # contains() is a literal substring match, so '%' or '_' in the query + # need no escaping (unlike LIKE/ILIKE). + condition = " OR ".join(f"contains(lower(coalesce({col}, '')), '{needle}')" for col in columns) + sql = ( + f"SELECT {id_column}, coalesce(title, '(no title)') AS title " + f"FROM {table} WHERE {condition} LIMIT {int(args.limit)}" + ) + try: + result = engine.run_query(con, sql, max_rows=0) + except duckdb.Error as exc: + print(f"\n{table.upper()}: search failed: {exc}", file=sys.stderr) + failures += 1 + continue + if result.rows: + suffix = " (showing first matches; raise --limit for more)" if len(result.rows) == args.limit else "" + print(f"\n{table.upper()}: {len(result.rows)} match{'es' if len(result.rows) != 1 else ''}{suffix}") + print("-" * 40) + for row_id, title in result.rows: + print(f" {row_id}: {str(title)[:80]}") + return 1 if failures else 0 diff --git a/src/spicy_regs/cli/stats.py b/src/spicy_regs/cli/stats.py new file mode 100644 index 0000000..901424d --- /dev/null +++ b/src/spicy_regs/cli/stats.py @@ -0,0 +1,57 @@ +"""``spicy-regs stats`` — row counts and top agencies for the core tables.""" + +from __future__ import annotations + +import argparse +import sys + +import duckdb + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir + +# The core record types (first entries of the published table list); the +# rollup/companion tables are better explored with `tables` + `query`. +CORE_TABLES = ("dockets", "documents", "comments") + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("stats", help="Show dataset statistics for the core tables") + add_source_arguments(parser) + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + + print("=" * 60) + print("Dataset Statistics") + print("=" * 60) + + failures = 0 + for table in CORE_TABLES: + spec = specs.get(table) + if spec is None: + print(f"\n{table.upper()}: Not downloaded yet (run: spicy-regs download)") + continue + con = engine.connect({table: spec}) + try: + total = engine.run_query(con, f"SELECT count(*) FROM {table}", max_rows=1).rows[0][0] + columns = [row[0] for row in engine.run_query(con, f"DESCRIBE {table}", max_rows=0).rows] + print(f"\n{table.upper()} ({spec.kind}: {spec.location})") + print("-" * 40) + print(f" Rows: {total:,}") + print(f" Columns: {', '.join(columns)}") + if "agency_code" in columns: + top = engine.run_query( + con, + f"SELECT agency_code, count(*) AS n FROM {table} GROUP BY 1 ORDER BY n DESC, agency_code LIMIT 5", + max_rows=5, + ) + print(" Top agencies:") + for agency, count in top.rows: + print(f" {agency}: {count:,}") + except duckdb.Error as exc: + print(f"\n{table.upper()}: could not read {spec.location}: {exc}", file=sys.stderr) + failures += 1 + return 1 if failures else 0 diff --git a/src/spicy_regs/cli/tables.py b/src/spicy_regs/cli/tables.py new file mode 100644 index 0000000..84515a8 --- /dev/null +++ b/src/spicy_regs/cli/tables.py @@ -0,0 +1,37 @@ +"""``spicy-regs tables`` — list the queryable tables and where each resolves.""" + +from __future__ import annotations + +import argparse +import json + +from spicy_regs.cli import engine +from spicy_regs.cli._common import add_source_arguments, get_output_dir + + +def register(subparsers: argparse._SubParsersAction) -> None: + parser = subparsers.add_parser("tables", help="List available tables and where each one resolves") + add_source_arguments(parser) + parser.add_argument("--format", choices=["table", "json"], default="table", help="Output format") + parser.set_defaults(run=run) + + +def run(args: argparse.Namespace) -> int: + specs = engine.resolve_view_specs(args.source, get_output_dir(args), args.r2_url) + if args.format == "json": + payload = [ + {"table": table, "source": spec.kind, "location": spec.location} for table, spec in sorted(specs.items()) + ] + print(json.dumps(payload, indent=2)) + return 0 + + if not specs: + print("No tables available for this source. Run: spicy-regs download") + return 1 + width = max(len(t) for t in specs) + for table, spec in sorted(specs.items()): + print(f"{table.ljust(width)} {spec.kind:5} {spec.location}") + missing = [t for t in engine.TABLES if t not in specs] + if missing: + print(f"\nNot available from this source: {', '.join(missing)}") + return 0 diff --git a/tests/test_cli_commands.py b/tests/test_cli_commands.py new file mode 100644 index 0000000..d9cb524 --- /dev/null +++ b/tests/test_cli_commands.py @@ -0,0 +1,146 @@ +"""End-to-end tests of the spicy-regs commands against local parquet fixtures.""" + +import csv +import io +import json +from pathlib import Path + +import pytest + +from spicy_regs.cli import main +from tests.conftest import COMMENT_SCHEMA, DOCKET_SCHEMA, DOCUMENT_SCHEMA, write_parquet_from_dicts + + +@pytest.fixture +def data_dir( + tmp_path: Path, sample_dockets: list[dict], sample_documents: list[dict], sample_comments: list[dict] +) -> Path: + write_parquet_from_dicts(tmp_path / "dockets.parquet", sample_dockets, DOCKET_SCHEMA) + write_parquet_from_dicts(tmp_path / "documents.parquet", sample_documents, DOCUMENT_SCHEMA) + write_parquet_from_dicts(tmp_path / "comments.parquet", sample_comments, COMMENT_SCHEMA) + return tmp_path + + +def run_cli(args: list[str]) -> int: + with pytest.raises(SystemExit) as excinfo: + main(args) + code = excinfo.value.code + assert isinstance(code, int) + return code + + +def test_query_table_format(data_dir: Path, capsys): + sql = "SELECT agency_code, count(*) AS n FROM dockets GROUP BY 1 ORDER BY 1" + assert run_cli(["query", sql, "--source", "local", "-o", str(data_dir)]) == 0 + out = capsys.readouterr().out + assert "agency_code" in out + assert "EPA" in out + assert "(2 rows)" in out + + +def test_query_json_format(data_dir: Path, capsys): + sql = "SELECT docket_id FROM dockets ORDER BY 1 LIMIT 1" + assert run_cli(["query", sql, "--source", "local", "-o", str(data_dir), "--format", "json"]) == 0 + assert json.loads(capsys.readouterr().out) == [{"docket_id": "EPA-2024-0001"}] + + +def test_query_joins_across_tables(data_dir: Path, capsys): + sql = ( + "SELECT d.docket_id, count(c.comment_id) AS n_comments " + "FROM dockets d JOIN comments c USING (docket_id) " + "GROUP BY 1 ORDER BY n_comments DESC, d.docket_id LIMIT 1" + ) + assert run_cli(["query", sql, "--source", "local", "-o", str(data_dir), "--format", "json"]) == 0 + assert json.loads(capsys.readouterr().out) == [{"docket_id": "EPA-2024-0001", "n_comments": 2}] + + +def test_query_csv_to_output_file(data_dir: Path, tmp_path: Path, capsys): + out_file = tmp_path / "result.csv" + sql = "SELECT docket_id, agency_code FROM dockets ORDER BY 1" + args = ["query", sql, "--source", "local", "-o", str(data_dir), "--format", "csv", "--output", str(out_file)] + assert run_cli(args) == 0 + assert f"Wrote 3 rows to {out_file}" in capsys.readouterr().out + rows = list(csv.reader(io.StringIO(out_file.read_text()))) + assert rows[0] == ["docket_id", "agency_code"] + assert len(rows) == 4 + + +def test_query_max_rows_truncates(data_dir: Path, capsys): + sql = "SELECT docket_id FROM dockets ORDER BY 1" + assert run_cli(["query", sql, "--source", "local", "-o", str(data_dir), "--max-rows", "1"]) == 0 + assert "more available" in capsys.readouterr().out + + +def test_query_bad_sql_fails_cleanly(data_dir: Path, capsys): + assert run_cli(["query", "SELECT nope FROM missing", "--source", "local", "-o", str(data_dir)]) == 1 + assert "Query failed" in capsys.readouterr().err + + +def test_query_empty_local_dir_errors(tmp_path: Path, capsys): + assert run_cli(["query", "SELECT 1", "--source", "local", "-o", str(tmp_path)]) == 1 + assert "No tables available" in capsys.readouterr().err + + +def test_tables_lists_local_and_missing(data_dir: Path, capsys): + assert run_cli(["tables", "--source", "local", "-o", str(data_dir)]) == 0 + out = capsys.readouterr().out + assert "dockets" in out + assert "Not available from this source" in out + assert "feed_summary" in out + + +def test_tables_json(data_dir: Path, capsys): + assert run_cli(["tables", "--source", "local", "-o", str(data_dir), "--format", "json"]) == 0 + payload = json.loads(capsys.readouterr().out) + assert {entry["table"] for entry in payload} == {"dockets", "documents", "comments"} + assert all(entry["source"] == "local" for entry in payload) + + +def test_describe_local_table(data_dir: Path, capsys): + assert run_cli(["describe", "dockets", "--source", "local", "-o", str(data_dir), "--format", "json"]) == 0 + columns = {row["column_name"] for row in json.loads(capsys.readouterr().out)} + assert {"docket_id", "agency_code", "title"} <= columns + + +def test_describe_unavailable_table(data_dir: Path, capsys): + assert run_cli(["describe", "feed_summary", "--source", "local", "-o", str(data_dir)]) == 1 + assert "not available locally" in capsys.readouterr().err + + +def test_stats(data_dir: Path, capsys): + assert run_cli(["stats", "--source", "local", "-o", str(data_dir)]) == 0 + out = capsys.readouterr().out + assert "DOCKETS" in out + assert "Rows: 3" in out + assert "EPA: 2" in out + + +def test_sample_with_agency_filter(data_dir: Path, capsys): + assert run_cli(["sample", "dockets", "-n", "10", "--agency", "EPA", "--source", "local", "-o", str(data_dir)]) == 0 + out = capsys.readouterr().out + assert "EPA-2024-0001" in out + assert "FDA-2024-0010" not in out + + +def test_search_hits_multiple_tables(data_dir: Path, capsys): + assert run_cli(["search", "drug", "--source", "local", "-o", str(data_dir)]) == 0 + out = capsys.readouterr().out + assert "DOCKETS" in out + assert "FDA-2024-0010" in out + + +def test_search_respects_limit(data_dir: Path, capsys): + assert run_cli(["search", "e", "--limit", "1", "--source", "local", "-o", str(data_dir)]) == 0 + assert "showing first matches" in capsys.readouterr().out + + +def test_agencies(data_dir: Path, capsys): + assert run_cli(["agencies", "--source", "local", "-o", str(data_dir)]) == 0 + out = capsys.readouterr().out + assert "EPA" in out + assert "FDA" in out + + +def test_no_command_prints_help(capsys): + assert run_cli([]) == 0 + assert "Available commands" in capsys.readouterr().out diff --git a/tests/test_cli_download.py b/tests/test_cli_download.py new file mode 100644 index 0000000..029635f --- /dev/null +++ b/tests/test_cli_download.py @@ -0,0 +1,104 @@ +"""Tests for `spicy-regs download` using httpx.MockTransport (no network).""" + +from pathlib import Path + +import httpx +import pytest + +from spicy_regs.cli import download + + +def make_client(files: dict[str, bytes]) -> httpx.Client: + """Client whose transport serves `files` keyed by URL path (e.g. '/dockets.parquet').""" + + def handler(request: httpx.Request) -> httpx.Response: + content = files.get(request.url.path) + if content is None: + return httpx.Response(404) + if request.method == "HEAD": + return httpx.Response(200, headers={"content-length": str(len(content))}) + return httpx.Response(200, content=content) + + return httpx.Client(transport=httpx.MockTransport(handler)) + + +def test_download_writes_atomically(tmp_path: Path): + client = make_client({"/dockets.parquet": b"PARQUET-BYTES"}) + dest = tmp_path / "dockets.parquet" + assert download.download_file(client, "https://bucket.test/dockets.parquet", dest) == "downloaded" + assert dest.read_bytes() == b"PARQUET-BYTES" + assert not list(tmp_path.glob("*.tmp")) + + +def test_download_skips_when_size_matches(tmp_path: Path, capsys): + client = make_client({"/dockets.parquet": b"12345"}) + dest = tmp_path / "dockets.parquet" + dest.write_bytes(b"12345") + assert download.download_file(client, "https://bucket.test/dockets.parquet", dest) == "skipped" + assert "up to date" in capsys.readouterr().out + + +def test_download_refreshes_when_size_differs(tmp_path: Path): + client = make_client({"/dockets.parquet": b"new-longer-content"}) + dest = tmp_path / "dockets.parquet" + dest.write_bytes(b"stale") + assert download.download_file(client, "https://bucket.test/dockets.parquet", dest) == "downloaded" + assert dest.read_bytes() == b"new-longer-content" + + +def test_force_redownloads_even_when_current(tmp_path: Path): + client = make_client({"/dockets.parquet": b"12345"}) + dest = tmp_path / "dockets.parquet" + dest.write_bytes(b"12345") + assert download.download_file(client, "https://bucket.test/dockets.parquet", dest, force=True) == "downloaded" + + +def test_unpublished_table_is_missing_not_failed(tmp_path: Path, capsys): + client = make_client({}) + dest = tmp_path / "court_dockets.parquet" + assert download.download_file(client, "https://bucket.test/court_dockets.parquet", dest) == "missing" + assert not dest.exists() + assert "not published" in capsys.readouterr().out + + +def test_network_failure_leaves_existing_file_intact(tmp_path: Path): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("boom") + + client = httpx.Client(transport=httpx.MockTransport(handler)) + dest = tmp_path / "dockets.parquet" + dest.write_bytes(b"precious") + assert download.download_file(client, "https://bucket.test/dockets.parquet", dest) == "failed" + assert dest.read_bytes() == b"precious" + assert not list(tmp_path.glob("*.tmp")) + + +def test_run_end_to_end_via_main(tmp_path: Path, monkeypatch, capsys): + files: dict[str, bytes] = { + f"/{name}.parquet": f"data-{name}".encode() for name in ("dockets", "documents", "comments") + } + monkeypatch.setattr(download, "_build_client", lambda transport=None: make_client(files)) + + from spicy_regs.cli import main + + with pytest.raises(SystemExit) as excinfo: + main(["download", "-o", str(tmp_path)]) + assert excinfo.value.code == 0 + for name in ("dockets", "documents", "comments"): + assert (tmp_path / f"{name}.parquet").read_bytes() == f"data-{name}".encode() + assert "Done!" in capsys.readouterr().out + + +def test_run_reports_failure_exit_code(tmp_path: Path, monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("boom") + + monkeypatch.setattr( + download, "_build_client", lambda transport=None: httpx.Client(transport=httpx.MockTransport(handler)) + ) + + from spicy_regs.cli import main + + with pytest.raises(SystemExit) as excinfo: + main(["download", "--tables", "dockets", "-o", str(tmp_path)]) + assert excinfo.value.code == 1 diff --git a/tests/test_cli_engine.py b/tests/test_cli_engine.py new file mode 100644 index 0000000..0f76acd --- /dev/null +++ b/tests/test_cli_engine.py @@ -0,0 +1,110 @@ +"""Tests for the CLI's DuckDB query engine (offline except the integration test).""" + +from datetime import datetime +from decimal import Decimal +from pathlib import Path + +import pytest + +from spicy_regs.cli import engine +from spicy_regs.cli._output import jsonify +from tests.conftest import COMMENT_SCHEMA, DOCKET_SCHEMA, write_parquet_from_dicts + + +@pytest.fixture +def data_dir(tmp_path: Path, sample_dockets: list[dict]) -> Path: + write_parquet_from_dicts(tmp_path / "dockets.parquet", sample_dockets, DOCKET_SCHEMA) + return tmp_path + + +def test_local_view_specs_only_lists_present_files(data_dir: Path): + specs = engine.local_view_specs(data_dir) + assert set(specs) == {"dockets"} + assert specs["dockets"].kind == "local" + assert specs["dockets"].location == str(data_dir / "dockets.parquet") + + +def test_local_view_specs_partitioned_comments(data_dir: Path, sample_comments: list[dict]): + partition = data_dir / "comments" / "agency_code=EPA" + partition.mkdir(parents=True) + records = [{k: v for k, v in c.items() if k != "agency_code"} for c in sample_comments if c["agency_code"] == "EPA"] + schema = {k: v for k, v in COMMENT_SCHEMA.items() if k != "agency_code"} + write_parquet_from_dicts(partition / "part-0.parquet", records, schema) + + specs = engine.local_view_specs(data_dir) + assert set(specs) == {"dockets", "comments"} + assert "hive_partitioning=true" in specs["comments"].sql + + con = engine.connect({"comments": specs["comments"]}) + result = engine.run_query(con, "SELECT DISTINCT agency_code FROM comments", max_rows=0) + assert result.rows == [("EPA",)] + + +def test_remote_view_specs_cover_all_tables(): + specs = engine.remote_view_specs("https://example.com/base/") + assert set(specs) == set(engine.TABLES) + assert specs["dockets"].location == "https://example.com/base/dockets.parquet" + assert specs["dockets"].kind == "r2" + + +def test_resolve_auto_prefers_local_and_falls_back_to_remote(data_dir: Path): + specs = engine.resolve_view_specs("auto", data_dir, "https://example.com") + assert set(specs) == set(engine.TABLES) + assert specs["dockets"].kind == "local" + assert all(specs[t].kind == "r2" for t in engine.TABLES if t != "dockets") + + +def test_resolve_rejects_unknown_source(tmp_path: Path): + with pytest.raises(ValueError, match="Unknown source"): + engine.resolve_view_specs("ftp", tmp_path, "https://example.com") + + +def test_run_query_truncation(data_dir: Path): + con = engine.connect(engine.local_view_specs(data_dir)) + truncated = engine.run_query(con, "SELECT docket_id FROM dockets ORDER BY 1", max_rows=2) + assert len(truncated.rows) == 2 + assert truncated.truncated is True + + exact = engine.run_query(con, "SELECT docket_id FROM dockets ORDER BY 1", max_rows=3) + assert len(exact.rows) == 3 + assert exact.truncated is False + + unlimited = engine.run_query(con, "SELECT docket_id FROM dockets ORDER BY 1", max_rows=0) + assert len(unlimited.rows) == 3 + assert unlimited.truncated is False + + +def test_connect_skips_unreadable_view(data_dir: Path, capsys): + missing = engine.ViewSpec( + table="documents", + kind="local", + location=str(data_dir / "documents.parquet"), + sql=f"SELECT * FROM read_parquet('{data_dir / 'documents.parquet'}')", + ) + con = engine.connect({**engine.local_view_specs(data_dir), "documents": missing}) + assert "warning: table documents not available" in capsys.readouterr().err + # The healthy view still works. + result = engine.run_query(con, "SELECT count(*) FROM dockets", max_rows=1) + assert result.rows == [(3,)] + + +def test_escape_sql_string(): + assert engine.escape_sql_string("O'Brien") == "O''Brien" + + +def test_jsonify_coercions(): + assert jsonify(datetime(2024, 6, 1, 12, 30)) == "2024-06-01T12:30:00" + assert jsonify(Decimal("1.50")) == "1.50" + assert jsonify(b"\x01\xff") == "01ff" + assert jsonify([1, Decimal("2")]) == [1, "2"] + assert jsonify({"k": datetime(2024, 1, 1)}) == {"k": "2024-01-01T00:00:00"} + assert jsonify(None) is None + + +@pytest.mark.integration +def test_remote_describe_against_live_r2(tmp_path: Path): + specs = engine.resolve_view_specs("r2", tmp_path) + con = engine.connect({"dockets": specs["dockets"]}) + result = engine.run_query(con, "DESCRIBE dockets", max_rows=0) + column_names = {row[0] for row in result.rows} + assert {"docket_id", "agency_code", "title"} <= column_names diff --git a/tests/test_cli_parser.py b/tests/test_cli_parser.py new file mode 100644 index 0000000..cf08218 --- /dev/null +++ b/tests/test_cli_parser.py @@ -0,0 +1,85 @@ +"""Argument-parsing tests for the spicy-regs CLI (no I/O).""" + +import pytest + +from spicy_regs.cli import build_parser +from spicy_regs.cli._registry import COMMANDS +from spicy_regs.data_dictionary import TABLES + + +@pytest.fixture +def parser(): + return build_parser() + + +def test_every_registered_command_parses(parser): + args_by_command = { + "download": [], + "tables": [], + "describe": ["dockets"], + "query": ["SELECT 1"], + "stats": [], + "sample": ["dockets"], + "search": ["climate"], + "agencies": [], + } + registered = {module.__name__.rsplit(".", 1)[-1] for module in COMMANDS} + assert registered == set(args_by_command) + for command, extra in args_by_command.items(): + args = parser.parse_args([command, *extra]) + assert args.command == command + assert callable(args.run) + + +def test_download_types_alias_still_parses(parser): + args = parser.parse_args(["download", "--types", "comments"]) + assert args.tables == ["comments"] + + +def test_download_tables_all_force(parser): + args = parser.parse_args(["download", "--tables", "dockets", "feed_summary", "--force"]) + assert args.tables == ["dockets", "feed_summary"] + assert args.force is True + assert args.download_all is False + + args = parser.parse_args(["download", "--all"]) + assert args.download_all is True + assert args.tables is None + + +def test_download_rejects_unknown_table(parser): + with pytest.raises(SystemExit): + parser.parse_args(["download", "--tables", "nonsense"]) + + +def test_query_defaults(parser): + args = parser.parse_args(["query", "SELECT 1"]) + assert args.sql == "SELECT 1" + assert args.source == "auto" + assert args.format == "table" + assert args.max_rows == 25 + assert args.output is None + + +def test_describe_rejects_unknown_table(parser): + with pytest.raises(SystemExit): + parser.parse_args(["describe", "not_a_table"]) + + +def test_sample_accepts_every_published_table(parser): + for table in TABLES: + args = parser.parse_args(["sample", table, "-n", "3"]) + assert args.data_type == table + assert args.n == 3 + + +def test_output_dir_accepted_before_and_after_subcommand(parser): + before = parser.parse_args(["-o", "/tmp/a", "stats"]) + assert before.output_dir == "/tmp/a" + after = parser.parse_args(["stats", "-o", "/tmp/b"]) + assert after.output_dir == "/tmp/b" + # The subcommand's flag must not clobber a value given before it. + both = parser.parse_args(["-o", "/tmp/a", "stats", "-o", "/tmp/b"]) + assert both.output_dir == "/tmp/b" + neither = parser.parse_args(["stats"]) + assert neither.output_dir is None