jp-idwr-db 0.2.2__py3-none-any.whl

jp_idwr_db/__init__.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from .api import get_data, get_latest_week, list_diseases, list_prefectures
+from .config import Config, configure, get_config
+from .data_manager import ensure_data
+from .datasets import load_dataset as load
+from .transform import merge, pivot
+from .types import DatasetName
+from .utils import attach_prefecture_id, prefecture_map
+
+__all__ = [
+    "Config",
+    "DatasetName",
+    "attach_prefecture_id",
+    "configure",
+    "ensure_data",
+    "get_config",
+    "get_data",
+    "get_latest_week",
+    "list_diseases",
+    "list_prefectures",
+    "load",
+    "merge",
+    "pivot",
+    "prefecture_map",
+]
+
+__version__ = "0.2.2"
+__data_version__ = __version__

jp_idwr_db/__main__.py

@@ -0,0 +1,7 @@
+"""Module entrypoint for ``python -m jp_idwr_db``."""
+
+from __future__ import annotations
+
+from .cli import main
+
+raise SystemExit(main())

jp_idwr_db/_internal/__init__.py

@@ -0,0 +1,11 @@
+"""Internal utilities for jp_idwr_db package.
+
+This module provides low-level functions for downloading and reading data files.
+These functions are used by the build pipeline and are not part of the public API.
+"""
+
+from __future__ import annotations
+
+from . import download, read, validation
+
+__all__ = ["download", "read", "validation"]

jp_idwr_db/_internal/download.py

@@ -0,0 +1,12 @@
+"""Internal download utilities.
+
+This module provides low-level download functionality.
+For simple data access, use the public jp.get_data() API instead.
+"""
+
+from __future__ import annotations
+
+# Re-export download functions from io.py
+from ..io import download, download_recent
+
+__all__ = ["download", "download_recent"]

jp_idwr_db/_internal/read.py

@@ -0,0 +1,12 @@
+"""Internal read utilities.
+
+This module provides low-level file reading functionality.
+For simple data access, use the public jp.get_data() API instead.
+"""
+
+from __future__ import annotations
+
+# Re-export read function from io.py
+from ..io import read
+
+__all__ = ["read"]

jp_idwr_db/_internal/validation.py

@@ -0,0 +1,133 @@
+"""Validation utilities for data quality checks.
+
+This module provides functions for validating data schemas, detecting duplicates,
+and ensuring data quality across different surveillance data sources.
+"""
+
+from __future__ import annotations
+
+from typing import cast
+
+import polars as pl
+
+
+def get_sentinel_only_diseases() -> set[str]:
+    """Get sentinel-only diseases (deprecated static helper).
+
+    Returns:
+        Empty set. Sentinel-only detection is now computed dynamically in
+        smart_merge() based on disease overlap with zensu data.
+    """
+    return set()
+
+
+def validate_schema(df: pl.DataFrame, required_columns: list[str] | None = None) -> None:
+    """Validate that a DataFrame has the required schema.
+
+    Args:
+        df: DataFrame to validate.
+        required_columns: List of required column names. If None, uses standard schema.
+
+    Raises:
+        ValueError: If required columns are missing.
+    """
+    if required_columns is None:
+        required_columns = ["prefecture", "year", "week", "disease", "count"]
+
+    missing = [col for col in required_columns if col not in df.columns]
+    if missing:
+        raise ValueError(f"Missing required columns: {missing}")
+
+
+def validate_no_duplicates(
+    df: pl.DataFrame,
+    keys: list[str] | None = None,
+) -> None:
+    """Validate that there are no duplicate records based on key columns.
+
+    Args:
+        df: DataFrame to validate.
+        keys: List of column names that define uniqueness. If None, uses
+              ["prefecture", "year", "week", "disease", "category"].
+              Category is included because the same (prefecture, year, week, disease)
+              can have multiple categories (e.g., "male", "female", "total").
+
+    Raises:
+        ValueError: If duplicate records are found.
+    """
+    if keys is None:
+        # Include category if it exists, since same disease can have multiple categories
+        keys = ["prefecture", "year", "week", "disease"]
+        if "category" in df.columns:
+            keys.append("category")
+
+    # Count occurrences of each unique combination
+    dups = df.group_by(keys).agg(pl.len().alias("count")).filter(pl.col("count") > 1)
+
+    if dups.height > 0:
+        raise ValueError(
+            f"Found {dups.height} duplicate records. First few duplicates:\n{dups.head(5)}"
+        )
+
+
+def validate_date_ranges(df: pl.DataFrame) -> None:
+    """Validate that year and week values are reasonable.
+
+    Args:
+        df: DataFrame to validate.
+
+    Raises:
+        ValueError: If year or week values are out of expected ranges.
+    """
+    if "year" in df.columns:
+        years = df["year"]
+        min_year_raw, max_year_raw = years.min(), years.max()
+        min_year = cast(int, min_year_raw)
+        max_year = cast(int, max_year_raw)
+        if min_year < 1999 or max_year > 2030:
+            raise ValueError(f"Year values out of expected range: {min_year}-{max_year}")
+
+    if "week" in df.columns:
+        weeks = df["week"]
+        min_week_raw, max_week_raw = weeks.min(), weeks.max()
+        min_week = cast(int, min_week_raw)
+        max_week = cast(int, max_week_raw)
+        if min_week < 1 or max_week > 53:
+            raise ValueError(f"Week values out of valid range: {min_week}-{max_week}")
+
+
+def smart_merge(
+    zensu_df: pl.DataFrame,
+    teiten_df: pl.DataFrame,
+) -> pl.DataFrame:
+    """Merge zensu and teiten data, preferring confirmed (zensu) data.
+
+    This function implements the "prefer confirmed" strategy:
+    - Keep ALL zensu (confirmed case) data
+    - Add ONLY sentinel diseases that are absent from zensu
+    - This avoids duplication while preserving diseases only in sentinel surveillance
+
+    Args:
+        zensu_df: Confirmed case data (from zensu/bullet files).
+        teiten_df: Sentinel surveillance data (from teiten files).
+
+    Returns:
+        Merged DataFrame with no duplicate diseases.
+
+    Example:
+        >>> zensu = pl.DataFrame({"disease": ["Influenza", "Tuberculosis"], "count": [100, 10]})
+        >>> teiten = pl.DataFrame({"disease": ["Influenza", "RSV"], "count": [120, 50]})
+        >>> merged = smart_merge(zensu, teiten)
+        >>> # Result: Influenza from zensu + RSV from teiten
+    """
+    confirmed_diseases = (
+        zensu_df.select("disease").drop_nulls().unique().get_column("disease").to_list()
+    )
+
+    # Filter teiten to only include diseases not present in confirmed data.
+    teiten_filtered = teiten_df.filter(~pl.col("disease").is_in(confirmed_diseases))
+
+    # Combine zensu (all diseases) + teiten (sentinel-only diseases)
+    merged = pl.concat([zensu_df, teiten_filtered], how="diagonal_relaxed")
+
+    return merged

jp_idwr_db/api.py

@@ -0,0 +1,203 @@
+"""New unified API for accessing infectious disease surveillance data.
+
+This module provides the main user-facing API for jp_idwr_db, offering
+simple data access with flexible filtering capabilities.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+import polars as pl
+
+from .datasets import load_dataset
+
+logger = logging.getLogger(__name__)
+
+
+def get_data(
+    disease: str | list[str] | None = None,
+    prefecture: str | list[str] | None = None,
+    year: int | tuple[int, int] | None = None,
+    week: int | tuple[int, int] | None = None,
+    source: Literal["confirmed", "sentinel", "all"] = "all",
+) -> pl.DataFrame:
+    """Get infectious disease surveillance data with optional filtering.
+
+    This is the main entry point for accessing jp_idwr_db data. It loads all
+    available data (historical + recent, confirmed + sentinel) and applies
+    optional filters.
+
+    Args:
+        disease: Filter by disease name(s). Case-insensitive partial matching.
+            Examples: "Influenza", ["COVID-19", "Influenza"], "RS virus"
+        prefecture: Filter by prefecture name(s).
+            Examples: "Tokyo", ["Tokyo", "Osaka"]
+        year: Filter by single year or (start, end) range (inclusive).
+            Examples: 2024, (2020, 2024)
+        week: Filter by single week or (start, end) range (inclusive).
+            Examples: 10, (1, 52)
+        source: Data source filter.
+            - "confirmed": Only zensu (all-case reporting) data
+            - "sentinel": Only teiten (sentinel surveillance) data
+            - "all": Both sources (default)
+
+    Returns:
+        DataFrame with standardized schema containing:
+        - prefecture: Prefecture name
+        - year: ISO year
+        - week: ISO week
+        - date: Week start date
+        - disease: Disease name (normalized)
+        - count: Weekly case count
+        - per_sentinel: Per-sentinel rate (sentinel only, null for confirmed)
+        - source: "Confirmed cases" or "Sentinel surveillance"
+        - category: "total", "male", "female" (when available)
+
+    Examples:
+        >>> import jp_idwr_db as jp
+        >>> # Get all data
+        >>> df = jp.get_data()
+
+        >>> # Filter by disease
+        >>> flu = jp.get_data(disease="Influenza")
+
+        >>> # Multiple diseases, specific year
+        >>> df = jp.get_data(disease=["COVID-19", "Influenza"], year=2024)
+
+        >>> # Prefecture and year range
+        >>> tokyo = jp.get_data(prefecture="Tokyo", year=(2020, 2024))
+
+        >>> # Only sentinel data
+        >>> sentinel = jp.get_data(source="sentinel", year=2024)
+
+        >>> # Complex filtering
+        >>> df = jp.get_data(
+        ...     disease=["Influenza", "RS virus"],
+        ...     prefecture=["Tokyo", "Osaka"],
+        ...     year=(2023, 2025),
+        ...     source="all"
+        ... )
+    """
+    # Load unified dataset (cached locally, downloaded from releases on demand).
+    try:
+        df = load_dataset("unified")
+    except Exception:
+        logger.warning("Failed to load unified dataset, falling back to bullet dataset")
+        try:
+            df = load_dataset("bullet")
+        except Exception:
+            logger.warning("Failed to load bullet dataset, returning empty DataFrame")
+            df = pl.DataFrame()
+
+    if df.height == 0:
+        return df
+
+    # Apply filters
+    if source != "all" and "source" in df.columns:
+        source_map = {
+            "confirmed": ["Confirmed cases", "All-case reporting"],
+            "sentinel": "Sentinel surveillance",
+        }
+        if source in source_map:
+            target = source_map[source]
+            if isinstance(target, list):
+                df = df.filter(pl.col("source").is_in(target))
+            else:
+                df = df.filter(pl.col("source") == target)
+
+    if disease is not None:
+        diseases = [disease] if isinstance(disease, str) else disease
+        # Case-insensitive partial matching
+        disease_filter = pl.lit(False)
+        for d in diseases:
+            disease_filter = disease_filter | pl.col("disease").str.to_lowercase().str.contains(
+                d.lower()
+            )
+        df = df.filter(disease_filter)
+
+    if prefecture is not None:
+        prefectures = [prefecture] if isinstance(prefecture, str) else prefecture
+        df = df.filter(pl.col("prefecture").is_in(prefectures))
+
+    if year is not None:
+        if isinstance(year, tuple):
+            start_year, end_year = year
+            df = df.filter((pl.col("year") >= start_year) & (pl.col("year") <= end_year))
+        else:
+            df = df.filter(pl.col("year") == year)
+
+    if week is not None:
+        if isinstance(week, tuple):
+            start_week, end_week = week
+            df = df.filter((pl.col("week") >= start_week) & (pl.col("week") <= end_week))
+        else:
+            df = df.filter(pl.col("week") == week)
+
+    return df
+
+
+def list_diseases(source: Literal["confirmed", "sentinel", "all"] = "all") -> list[str]:
+    """Get list of available disease names.
+
+    Args:
+        source: Filter by data source - "confirmed", "sentinel", or "all".
+
+    Returns:
+        Sorted list of disease names.
+
+    Example:
+        >>> import jp_idwr_db as jp
+        >>> all_diseases = jp.list_diseases()
+        >>> sentinel_only = jp.list_diseases(source="sentinel")
+    """
+    df = get_data(source=source)
+    if df.height == 0:
+        return []
+    return sorted(df["disease"].unique().to_list())
+
+
+def list_prefectures() -> list[str]:
+    """Get list of prefecture names.
+
+    Returns:
+        Sorted list of prefecture names.
+
+    Example:
+        >>> import jp_idwr_db as jp
+        >>> prefectures = jp.list_prefectures()
+        >>> print(prefectures[:3])
+        ['Aichi', 'Akita', 'Aomori']
+    """
+    df = get_data()
+    if df.height == 0:
+        return []
+    return sorted(df["prefecture"].unique().to_list())
+
+
+def get_latest_week() -> tuple[int, int] | None:
+    """Get the latest (year, week) with data available.
+
+    Returns:
+        Tuple of (year, week) for the most recent data, or None if no data.
+
+    Example:
+        >>> import jp_idwr_db as jp
+        >>> latest = jp.get_latest_week()
+        >>> if latest:
+        ...     year, week = latest
+        ...     print(f"Latest data: {year} week {week}")
+    """
+    df = get_data()
+    if df.height == 0:
+        return None
+
+    # Check if year column exists, otherwise we can't determine the latest week
+    if "year" not in df.columns or "week" not in df.columns:
+        logger.warning("Cannot determine latest week: missing year or week column")
+        return None
+
+    # Get row with maximum year, then maximum week within that year
+    latest = df.sort(["year", "week"], descending=True).head(1)
+    return (int(latest["year"][0]), int(latest["week"][0]))

jp_idwr_db/cli.py

@@ -0,0 +1,38 @@
+"""Command-line interface for jp_idwr_db."""
+
+from __future__ import annotations
+
+import argparse
+
+from .data_manager import ensure_data
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the top-level CLI argument parser."""
+    parser = argparse.ArgumentParser(prog="jp-idwr-db")
+    subparsers = parser.add_subparsers(dest="command")
+
+    data_parser = subparsers.add_parser("data", help="Manage local data cache")
+    data_subparsers = data_parser.add_subparsers(dest="data_command")
+
+    download_parser = data_subparsers.add_parser("download", help="Download release parquet assets")
+    download_parser.add_argument(
+        "--version", type=str, default=None, help="Data version (e.g. v0.1.0)"
+    )
+    download_parser.add_argument("--force", action="store_true", help="Force re-download")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Run CLI entrypoint."""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "data" and args.data_command == "download":
+        data_dir = ensure_data(version=args.version, force=args.force)
+        print(data_dir)
+        return 0
+
+    parser.print_help()
+    return 0

jp_idwr_db/config.py

@@ -0,0 +1,60 @@
+"""Configuration management for jp_idwr_db.
+
+This module provides a global configuration system for controlling package behavior.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from pathlib import Path
+
+from platformdirs import user_cache_dir
+
+
+@dataclass(frozen=True)
+class Config:
+    """Global configuration for jp_idwr_db package.
+
+    Attributes:
+        cache_dir: Directory for caching downloaded files.
+        rate_limit_per_minute: Maximum number of HTTP requests per minute.
+        user_agent: User-Agent header for HTTP requests.
+        timeout_seconds: Timeout for HTTP requests in seconds.
+        retries: Number of retry attempts for failed requests.
+    """
+
+    cache_dir: Path = Path(user_cache_dir("jp_idwr_db"))
+    rate_limit_per_minute: int = 20
+    user_agent: str = "jp_idwr_db/0.2.2 (+https://github.com/AlFontal/jp-idwr-db)"
+    timeout_seconds: float = 30.0
+    retries: int = 3
+
+
+_CONFIG = Config()
+
+
+def get_config() -> Config:
+    """Get the current global configuration.
+
+    Returns:
+        The current Config instance.
+    """
+    return _CONFIG
+
+
+def configure(**kwargs: object) -> Config:
+    """Update the global configuration.
+
+    Args:
+        **kwargs: Configuration parameters to update (see Config attributes).
+
+    Returns:
+        The updated Config instance.
+
+    Example:
+        >>> import jp_idwr_db as jp
+        >>> jp.configure(rate_limit_per_minute=10)
+    """
+    global _CONFIG  # noqa: PLW0603
+    _CONFIG = replace(_CONFIG, **kwargs)  # type: ignore[arg-type]
+    return _CONFIG

jp_idwr_db/data_manager.py

@@ -0,0 +1,162 @@
+"""Runtime dataset manager for release-hosted parquet assets."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import zipfile
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as package_version
+from pathlib import Path
+from typing import Any
+
+import httpx
+from platformdirs import user_cache_dir
+
+PACKAGE_NAME = "jp_idwr_db"
+DEFAULT_REPO = "AlFontal/jp-idwr-db"
+DEFAULT_BASE_URL = f"https://github.com/{DEFAULT_REPO}/releases/download"
+ARCHIVE_NAME = "jp_idwr_db-parquet.zip"
+MANIFEST_NAME = "jp_idwr_db-manifest.json"
+EXPECTED_DATASETS = {
+    "sex_prefecture.parquet",
+    "place_prefecture.parquet",
+    "bullet.parquet",
+    "sentinel.parquet",
+    "unified.parquet",
+    "prefecture_en.parquet",
+}
+
+
+def get_cache_dir() -> Path:
+    """Return the base cache directory for package data."""
+    override = os.getenv("JPINFECT_CACHE_DIR")
+    if override:
+        return Path(override).expanduser()
+    return Path(user_cache_dir(PACKAGE_NAME))
+
+
+def _resolve_data_version(version: str | None) -> str:
+    """Resolve data version from explicit arg, env var, or package version."""
+    if version:
+        return version
+    env_version = os.getenv("JPINFECT_DATA_VERSION")
+    if env_version:
+        return env_version
+    try:
+        pkg_version = package_version("jp-idwr-db")
+    except PackageNotFoundError:
+        pkg_version = "0.0.0"
+    return pkg_version if pkg_version.startswith("v") else f"v{pkg_version}"
+
+
+def _resolve_base_url(version: str) -> str:
+    """Resolve base URL for release assets."""
+    base_url = os.getenv("JPINFECT_DATA_BASE_URL")
+    if base_url:
+        return base_url.rstrip("/")
+    return f"{DEFAULT_BASE_URL}/{version}"
+
+
+def _sha256(path: Path) -> str:
+    """Compute SHA256 hash for a file path."""
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _download_file(url: str, dest: Path) -> None:
+    """Download URL content to a local path."""
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    with httpx.stream("GET", url, timeout=60.0, follow_redirects=True) as response:
+        response.raise_for_status()
+        with dest.open("wb") as handle:
+            for chunk in response.iter_bytes():
+                handle.write(chunk)
+
+
+def _verify_manifest(manifest: dict[str, Any]) -> None:
+    """Ensure manifest has expected structure."""
+    required = {"archive", "archive_sha256", "files"}
+    missing = required - set(manifest)
+    if missing:
+        raise ValueError(f"Invalid manifest; missing keys: {sorted(missing)}")
+    files = manifest["files"]
+    if not isinstance(files, dict) or not files:
+        raise ValueError("Invalid manifest; 'files' must be a non-empty object")
+
+
+def download_release_assets(version: str, dest_dir: Path) -> tuple[Path, Path]:
+    """Download release archive + manifest for a data version."""
+    base_url = _resolve_base_url(version)
+    manifest_path = dest_dir / MANIFEST_NAME
+    archive_path = dest_dir / ARCHIVE_NAME
+
+    _download_file(f"{base_url}/{MANIFEST_NAME}", manifest_path)
+    _download_file(f"{base_url}/{ARCHIVE_NAME}", archive_path)
+    return archive_path, manifest_path
+
+
+def _extract_archive(archive_path: Path, dest_dir: Path) -> None:
+    """Extract archive into destination directory."""
+    with zipfile.ZipFile(archive_path) as archive:
+        archive.extractall(dest_dir)
+
+
+def ensure_data(version: str | None = None, force: bool = False) -> Path:
+    """Ensure parquet assets are available in the local cache.
+
+    Args:
+        version: Data release version (for example, ``v0.1.0``).
+        force: Re-download and replace cached files.
+
+    Returns:
+        Directory path containing parquet files for the resolved version.
+    """
+    resolved = _resolve_data_version(version)
+    cache_dir = get_cache_dir()
+    data_dir = cache_dir / "data" / resolved
+    marker = data_dir / ".complete"
+
+    if marker.exists() and not force:
+        return data_dir
+
+    if force and data_dir.exists():
+        shutil.rmtree(data_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    archive_path, manifest_path = download_release_assets(resolved, data_dir)
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    _verify_manifest(manifest)
+
+    expected_archive = str(manifest["archive"])
+    if archive_path.name != expected_archive:
+        expected_path = data_dir / expected_archive
+        archive_path.rename(expected_path)
+        archive_path = expected_path
+
+    archive_hash = _sha256(archive_path)
+    if archive_hash != manifest["archive_sha256"]:
+        raise ValueError("Archive checksum mismatch")
+
+    _extract_archive(archive_path, data_dir)
+
+    file_entries: dict[str, dict[str, Any]] = manifest["files"]
+    for rel_name, file_info in file_entries.items():
+        file_path = data_dir / rel_name
+        if not file_path.exists():
+            raise ValueError(f"Missing extracted data file: {rel_name}")
+        expected_hash = str(file_info["sha256"])
+        if _sha256(file_path) != expected_hash:
+            raise ValueError(f"Checksum mismatch for {rel_name}")
+
+    missing_expected = [name for name in EXPECTED_DATASETS if not (data_dir / name).exists()]
+    if missing_expected:
+        raise ValueError(f"Missing required datasets in cache: {sorted(missing_expected)}")
+
+    marker.write_text("ok\n", encoding="utf-8")
+    return data_dir