fixes and refactor

app/services/mortality.py

@@ -0,0 +1,127 @@
+"""
+Mortality/Life table utilities.
+
+Helpers to query `life_tables` and `number_tables` by age/month and
+return values filtered by sex/race using compact codes:
+  - sex: M, F, A (All)
+  - race: W (White), B (Black), H (Hispanic), A (All)
+
+Column naming in tables follows the pattern:
+  - LifeTable: le_{race}{sex}, na_{race}{sex}
+  - NumberTable: na_{race}{sex}
+
+Examples:
+  - race=W, sex=M => suffix "wm" (columns `le_wm`, `na_wm`)
+  - race=A, sex=F => suffix "af" (columns `le_af`, `na_af`)
+  - race=H, sex=A => suffix "ha" (columns `le_ha`, `na_ha`)
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Optional, Tuple
+from sqlalchemy.orm import Session
+
+from app.models.pensions import LifeTable, NumberTable
+
+
+_RACE_MAP: Dict[str, str] = {
+    "W": "w",  # White
+    "B": "b",  # Black
+    "H": "h",  # Hispanic
+    "A": "a",  # All races
+}
+
+_SEX_MAP: Dict[str, str] = {
+    "M": "m",
+    "F": "f",
+    "A": "a",  # All sexes
+}
+
+
+class InvalidCodeError(ValueError):
+    pass
+
+
+def _normalize_codes(sex: str, race: str) -> Tuple[str, str, str]:
+    """Validate/normalize sex and race to construct the column suffix.
+
+    Returns (suffix, sex_u, race_u) where suffix is lowercase like "wm".
+    Raises InvalidCodeError on invalid inputs.
+    """
+    sex_u = (sex or "").strip().upper()
+    race_u = (race or "").strip().upper()
+    if sex_u not in _SEX_MAP:
+        raise InvalidCodeError(f"Invalid sex code '{sex}'. Expected one of: {', '.join(_SEX_MAP.keys())}")
+    if race_u not in _RACE_MAP:
+        raise InvalidCodeError(f"Invalid race code '{race}'. Expected one of: {', '.join(_RACE_MAP.keys())}")
+    return _RACE_MAP[race_u] + _SEX_MAP[sex_u], sex_u, race_u
+
+
+def get_life_values(
+    db: Session,
+    *,
+    age: int,
+    sex: str,
+    race: str,
+) -> Optional[Dict[str, Optional[float]]]:
+    """Return life table LE and NA values for a given age, sex, and race.
+
+    Returns dict: {"age": int, "sex": str, "race": str, "le": float|None, "na": float|None}
+    Returns None if the age row does not exist.
+    Raises InvalidCodeError for invalid codes.
+    """
+    suffix, sex_u, race_u = _normalize_codes(sex, race)
+    row: Optional[LifeTable] = db.query(LifeTable).filter(LifeTable.age == age).first()
+    if not row:
+        return None
+
+    le_col = f"le_{suffix}"
+    na_col = f"na_{suffix}"
+    le_val = getattr(row, le_col, None)
+    na_val = getattr(row, na_col, None)
+
+    return {
+        "age": int(age),
+        "sex": sex_u,
+        "race": race_u,
+        "le": float(le_val) if le_val is not None else None,
+        "na": float(na_val) if na_val is not None else None,
+    }
+
+
+def get_number_value(
+    db: Session,
+    *,
+    month: int,
+    sex: str,
+    race: str,
+) -> Optional[Dict[str, Optional[float]]]:
+    """Return number table NA value for a given month, sex, and race.
+
+    Returns dict: {"month": int, "sex": str, "race": str, "na": float|None}
+    Returns None if the month row does not exist.
+    Raises InvalidCodeError for invalid codes.
+    """
+    suffix, sex_u, race_u = _normalize_codes(sex, race)
+    row: Optional[NumberTable] = db.query(NumberTable).filter(NumberTable.month == month).first()
+    if not row:
+        return None
+
+    na_col = f"na_{suffix}"
+    na_val = getattr(row, na_col, None)
+
+    return {
+        "month": int(month),
+        "sex": sex_u,
+        "race": race_u,
+        "na": float(na_val) if na_val is not None else None,
+    }
+
+
+__all__ = [
+    "InvalidCodeError",
+    "get_life_values",
+    "get_number_value",
+]
+
+