skills/genebass-gene-burden-skill/scripts/genebass_gene_burden.py

#!/usr/bin/env python3
"""genebass-gene-burden

Fetch Genebass gene burden PheWAS associations for one Ensembl gene input.
Input JSON on stdin:
  - {"ensembl_gene_id":"ENSG00000173531"}
  - {"ensembl_gene_id":"ENSG00000173531","burden_set":"pLoF","max_results":100}
  - "ENSG00000173531"
Output JSON on stdout.
"""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path
from typing import Any
from urllib.parse import quote, unquote

import requests

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

GENEBASS_API_BASE = "https://main.genebass.org/api"
USER_AGENT = "genebass-gene-burden-skill/1.0 (+requests)"
DEFAULT_TIMEOUT_S = 30
CANONICAL_BURDEN_SETS = ("pLoF", "missense|LC", "synonymous")
ENSG_RE = re.compile(r"^ENSG[0-9]+(?:\.[0-9]+)?$", re.IGNORECASE)


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 normalize_gene_id(raw: str) -> str:
    value = raw.strip()
    if not value:
        raise ValueError("Ensembl gene ID is empty.")
    if not ENSG_RE.match(value):
        raise ValueError(
            "Invalid Ensembl gene ID. Expected format like ENSG00000173531 (optional .version)."
        )
    value = value.upper()
    if "." in value:
        value = value.split(".", 1)[0]
    return value


def normalize_burden_set(raw: str) -> str:
    value = unquote(raw.strip())
    key = re.sub(r"[^a-z0-9]", "", value.lower())

    if key in {"plof", "lof"}:
        return "pLoF"
    if key in {"missense", "missenselc"}:
        return "missense|LC"
    if key == "synonymous":
        return "synonymous"

    canonical = ", ".join(CANONICAL_BURDEN_SETS)
    raise ValueError(
        f"Invalid `burden_set`. Allowed canonical values: {canonical}. "
        "Accepted aliases: LoF/LOF/lof/plof -> pLoF; missense -> missense|LC."
    )


def parse_input(payload: Any) -> tuple[str, str, int | None]:
    if isinstance(payload, str):
        return normalize_gene_id(payload), "pLoF", None

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

    gene = payload.get("ensembl_gene_id") or payload.get("gene_id") or payload.get("gene")
    if not gene or not isinstance(gene, str):
        raise ValueError("Provide `ensembl_gene_id` as a non-empty string.")
    gene_id = normalize_gene_id(gene)

    burden_raw = payload.get("burden_set", "pLoF")
    if not isinstance(burden_raw, str) or not burden_raw.strip():
        raise ValueError("`burden_set` must be a non-empty string when provided.")
    burden_set = normalize_burden_set(burden_raw)

    max_results = payload.get("max_results")
    if max_results is None:
        return gene_id, burden_set, None

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


def fetch_gene_phewas(
    gene_id: str, burden_set: str, *, session: requests.Session | None = None
) -> tuple[str, Any, int]:
    encoded_gene = quote(gene_id, safe="")
    encoded_burden = quote(burden_set, safe="|")
    url = f"{GENEBASS_API_BASE}/phewas/{encoded_gene}?burdenSet={encoded_burden}"
    headers = {
        "Accept": "application/json",
        "User-Agent": USER_AGENT,
    }

    requester = session if session is not None else requests
    resp = requester.get(url, headers=headers, timeout=DEFAULT_TIMEOUT_S)
    status_code = resp.status_code
    if status_code == 404:
        return url, None, status_code
    if status_code >= 400:
        raise RuntimeError(f"Genebass API returned HTTP {status_code}.")
    try:
        data = resp.json()
    except ValueError as exc:
        raise RuntimeError("Genebass API returned non-JSON.") from exc
    return url, data, status_code


def fetch_phenotypes_metadata(*, session: requests.Session | None = None) -> list[dict[str, Any]]:
    url = f"{GENEBASS_API_BASE}/phenotypes"
    headers = {
        "Accept": "application/json",
        "User-Agent": USER_AGENT,
    }

    requester = session if session is not None else requests
    resp = requester.get(url, headers=headers, timeout=DEFAULT_TIMEOUT_S)
    if resp.status_code >= 400:
        raise RuntimeError(f"Genebass phenotypes API returned HTTP {resp.status_code}.")
    try:
        data = resp.json()
    except ValueError as exc:
        raise RuntimeError("Genebass phenotypes API returned non-JSON.") from exc

    if isinstance(data, list):
        return [row for row in data if isinstance(row, dict)]
    return []


def build_description_map(phenotypes: list[dict[str, Any]]) -> dict[str, str]:
    mapping: dict[str, str] = {}
    for row in phenotypes:
        analysis_id = row.get("analysis_id")
        description = row.get("description")
        if isinstance(analysis_id, str) and analysis_id and isinstance(description, str):
            mapping[analysis_id] = description
    return mapping


def unpack_phewas_payload(
    data: Any,
) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
    if data is None:
        return None, []

    if isinstance(data, dict):
        gene = data.get("gene") if isinstance(data.get("gene"), dict) else None
        rows = data.get("phewas") if isinstance(data.get("phewas"), list) else []
        return gene, [row for row in rows if isinstance(row, dict)]

    if isinstance(data, list) and data:
        first = data[0]
        if isinstance(first, dict) and isinstance(first.get("phewas"), list):
            gene = first.get("gene") if isinstance(first.get("gene"), dict) else None
            rows = first.get("phewas") or []
            return gene, [row for row in rows if isinstance(row, dict)]
        return None, [row for row in data if isinstance(row, dict)]

    return None, []


def build_phenotype_id(row: dict[str, Any]) -> str:
    trait_type = row.get("trait_type", "")
    phenocode = row.get("phenocode", "")
    pheno_sex = row.get("pheno_sex", "")
    coding = row.get("coding", "")
    modifier = row.get("modifier", "")
    return f"{trait_type}-{phenocode}-{pheno_sex}-{coding}-{modifier}"


def transform_rows(
    rows: list[dict[str, Any]], description_map: dict[str, str]
) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    for row in rows:
        phenotype_id = build_phenotype_id(row)
        out.append(
            {
                "phenotype_id": phenotype_id,
                "phenotype_description": description_map.get(phenotype_id),
                "skat_o_pvalue": row.get("Pvalue"),
            }
        )
    return out


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:
        gene_id, burden_set, max_results = parse_input(payload)
    except ValueError as exc:
        sys.stdout.write(json.dumps(error("invalid_input", str(exc))))
        return 2

    with requests.Session() as session:
        try:
            query_url, data, status_code = fetch_gene_phewas(gene_id, burden_set, session=session)
        except requests.RequestException as exc:
            sys.stdout.write(
                json.dumps(
                    error(
                        "network_error",
                        f"Genebass request failed ({type(exc).__name__}).",
                    )
                )
            )
            return 1
        except RuntimeError as exc:
            message = str(exc)
            code = "upstream_error"
            if "HTTP 500" in message:
                code = "not_found_or_invalid_upstream_request"
                message = (
                    "Genebass returned HTTP 500. This often means an unknown gene ID or invalid "
                    "burden_set. Use Ensembl IDs and burden_set in "
                    "{pLoF, missense|LC, synonymous} (aliases are accepted)."
                )
            sys.stdout.write(json.dumps(error(code, message)))
            return 1

        if status_code == 404:
            warnings.append("Gene not found in Genebass PheWAS API (HTTP 404).")
            output = {
                "ok": True,
                "source": "genebass",
                "input": {"type": "ensembl_gene_id", "value": gene_id},
                "burden_set": burden_set,
                "query_url": query_url,
                "gene": None,
                "association_count": 0,
                "association_count_total": 0,
                "truncated": False,
                "associations": [],
                "warnings": warnings,
            }
            output = apply_source_contract(
                output, "genebass-gene-burden-skill", query_url, mode="empty"
            )
            sys.stdout.write(json.dumps(output))
            return 0

        description_map: dict[str, str] = {}
        try:
            phenotypes = fetch_phenotypes_metadata(session=session)
            description_map = build_description_map(phenotypes)
        except requests.RequestException as exc:
            warnings.append(f"Could not fetch phenotype descriptions ({type(exc).__name__}).")
        except RuntimeError as exc:
            warnings.append(str(exc))

    gene, rows = unpack_phewas_payload(data)
    associations = transform_rows(rows, description_map)
    total = len(associations)
    if max_results is not None and total > max_results:
        associations = associations[:max_results]
    truncated = len(associations) < total

    gene_out = None
    if isinstance(gene, dict):
        gene_out = {
            "gene_id": gene.get("gene_id"),
            "symbol": gene.get("symbol"),
            "name": gene.get("name"),
        }

    output = {
        "ok": True,
        "source": "genebass",
        "input": {"type": "ensembl_gene_id", "value": gene_id},
        "burden_set": burden_set,
        "query_url": query_url,
        "gene": gene_out,
        "association_count": len(associations),
        "association_count_total": total,
        "truncated": truncated,
        "associations": associations,
        "warnings": warnings,
    }
    output = apply_source_contract(output, "genebass-gene-burden-skill", query_url)
    sys.stdout.write(json.dumps(output))
    return 0


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