skills/tpmi-phewas-skill/scripts/tpmi_phewas.py

#!/usr/bin/env python3
"""tpmi-phewas

Fetch TPMI PheWAS associations for one variant input.
Input JSON on stdin:
  - {"grch38":"6:160540105-T-C"}
  - {"grch37":"6:162447146:T:C","max_results":25}
  - {"rsid":"rs9273363","max_results":25,"save_raw":true}
  - "6:160540105-T-C"
Output JSON on stdout.
"""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path
from typing import Any

import requests
from variant_resolution import (
    VariantResolutionError,
    extract_variant_input,
    resolve_query_variant,
)

sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "scripts"))
from database_source_contract import apply_source_contract  # noqa: E402

TPMI_BASE = "https://pheweb.ibms.sinica.edu.tw"
USER_AGENT = "tpmi-phewas-skill/1.0 (+requests)"
DEFAULT_TIMEOUT_S = 20
DEFAULT_MAX_RESULTS = 10
SAFE_PATH_RE = re.compile(r"[^A-Za-z0-9._-]+")


def error(code: str, message: str, warnings: list[str] | None = None) -> dict[str, Any]:
    return {
        "ok": False,
        "error": {"code": code, "message": message},
        "warnings": warnings or [],
    }


def parse_input(payload: Any) -> tuple[str, str, int, bool, str | None, float]:
    if isinstance(payload, str):
        return (
            "grch38",
            payload.strip(),
            DEFAULT_MAX_RESULTS,
            False,
            None,
            DEFAULT_TIMEOUT_S,
        )

    if not isinstance(payload, dict):
        raise ValueError("Input must be a JSON string or object.")

    input_type, variant = extract_variant_input(payload, default_build_key="grch38")

    max_results = payload.get("max_results", DEFAULT_MAX_RESULTS)
    if not isinstance(max_results, int) or max_results <= 0:
        raise ValueError("`max_results` must be a positive integer when provided.")

    save_raw = payload.get("save_raw", False)
    if not isinstance(save_raw, bool):
        raise ValueError("`save_raw` must be a boolean when provided.")

    raw_output_path = payload.get("raw_output_path")
    if raw_output_path is not None:
        if not isinstance(raw_output_path, str) or not raw_output_path.strip():
            raise ValueError("`raw_output_path` must be a non-empty string when provided.")
        raw_output_path = raw_output_path.strip()

    timeout_sec = payload.get("timeout_sec", DEFAULT_TIMEOUT_S)
    if not isinstance(timeout_sec, (int, float)) or timeout_sec <= 0:
        raise ValueError("`timeout_sec` must be a positive number when provided.")

    return (
        input_type,
        variant,
        max_results,
        save_raw,
        raw_output_path,
        float(timeout_sec),
    )


def fetch_tpmi_variant(
    session: requests.Session,
    variant_str: str,
    timeout_sec: float,
) -> tuple[Any | None, int | None, bytes | None]:
    encoded = requests.utils.quote(variant_str, safe=":-")
    url = f"{TPMI_BASE}/api/variant/{encoded}"
    headers = {
        "Accept": "application/json",
        "User-Agent": USER_AGENT,
    }

    resp = session.get(url, headers=headers, timeout=timeout_sec)
    if resp.status_code == 404:
        return None, 404, None
    resp.raise_for_status()
    raw_content = getattr(resp, "content", None)
    raw_response = (
        bytes(raw_content)
        if isinstance(raw_content, (bytes, bytearray, memoryview))
        else resp.text.encode("utf-8")
    )
    return resp.json(), resp.status_code, raw_response


def extract_associations(data: Any) -> list[Any]:
    if data is None:
        return []
    if isinstance(data, dict) and isinstance(data.get("phenos"), list):
        return data["phenos"]
    if isinstance(data, dict) and isinstance(data.get("results"), list):
        return data["results"]
    if isinstance(data, list):
        return data
    return []


def resolve_raw_output_path(canonical_variant: str, raw_output_path: str | None) -> Path:
    if raw_output_path:
        return Path(raw_output_path).expanduser()

    safe_variant = SAFE_PATH_RE.sub("_", canonical_variant).strip("._") or "variant"
    return Path("/tmp") / f"tpmi-phewas-{safe_variant}.json"


def write_raw_json(path: Path, raw_response: bytes | str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    raw_bytes = raw_response if isinstance(raw_response, bytes) else raw_response.encode("utf-8")
    path.write_bytes(raw_bytes)


def main() -> int:
    warnings: list[str] = []

    try:
        payload = json.load(sys.stdin)
    except Exception as exc:  # noqa: BLE001
        sys.stdout.write(json.dumps(error("invalid_json", f"Could not parse JSON input: {exc}")))
        return 2

    try:
        (
            input_type,
            input_variant,
            max_results,
            save_raw,
            raw_output_path,
            timeout_sec,
        ) = parse_input(payload)
    except ValueError as exc:
        sys.stdout.write(json.dumps(error("invalid_input", str(exc))))
        return 2

    try:
        resolution = resolve_query_variant(
            input_type=input_type,
            input_value=input_variant,
            target_build="GRCh38",
        )
        parsed = dict(resolution["query_variant"])
        warnings.extend(resolution["warnings"])
    except VariantResolutionError as exc:
        sys.stdout.write(json.dumps(error(exc.code, exc.message, exc.warnings)))
        return 1
    except requests.RequestException as exc:
        sys.stdout.write(
            json.dumps(
                error(
                    "network_error",
                    f"Variant resolution failed ({type(exc).__name__}).",
                )
            )
        )
        return 1

    session = requests.Session()
    try:
        data, status_code, raw_response = fetch_tpmi_variant(
            session, parsed["canonical"], timeout_sec
        )
    except requests.RequestException as exc:
        sys.stdout.write(
            json.dumps(error("network_error", f"TPMI request failed ({type(exc).__name__})."))
        )
        return 1
    except ValueError as exc:
        sys.stdout.write(json.dumps(error("invalid_response", f"TPMI returned non-JSON: {exc}")))
        return 1

    variant_url = f"{TPMI_BASE}/variant/{parsed['canonical']}"
    saved_raw_output_path: str | None = None
    if save_raw and data is not None:
        raw_path = resolve_raw_output_path(parsed["canonical"], raw_output_path)
        try:
            write_raw_json(raw_path, raw_response or b"")
        except OSError as exc:
            sys.stdout.write(json.dumps(error("write_error", f"Could not write raw output: {exc}")))
            return 1
        saved_raw_output_path = str(raw_path)

    if status_code == 404 or data is None:
        warnings.append("Variant not found in TPMI PheWAS API.")
        output = {
            "ok": True,
            "source": "tpmi",
            "input": resolution["input"],
            "query_variant": parsed,
            "max_results_applied": max_results,
            "association_count": 0,
            "association_count_total": 0,
            "truncated": False,
            "associations": [],
            "variant": None,
            "variant_url": variant_url,
            "raw_output_path": None,
            "warnings": warnings,
        }
        output = apply_source_contract(output, "tpmi-phewas-skill", variant_url, mode="empty")
        sys.stdout.write(json.dumps(output))
        return 0

    associations = extract_associations(data)
    total = len(associations)
    if total > max_results:
        associations = associations[:max_results]
    truncated = len(associations) < total

    variant_info = None
    if isinstance(data, dict):
        variant_info = {
            "chrom": data.get("chrom"),
            "pos": data.get("pos"),
            "ref": data.get("ref"),
            "alt": data.get("alt"),
            "rsids": data.get("rsids"),
            "variant_name": data.get("variant_name"),
            "nearest_genes": data.get("nearest_genes"),
        }

    output = {
        "ok": True,
        "source": "tpmi",
        "input": resolution["input"],
        "query_variant": parsed,
        "max_results_applied": max_results,
        "association_count": len(associations),
        "association_count_total": total,
        "truncated": truncated,
        "associations": associations,
        "variant": variant_info,
        "variant_url": variant_url,
        "raw_output_path": saved_raw_output_path,
        "warnings": warnings,
    }
    output = apply_source_contract(output, "tpmi-phewas-skill", variant_url)
    sys.stdout.write(json.dumps(output))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())