128 lines
3.4 KiB
Python
128 lines
3.4 KiB
Python
"""
|
|
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",
|
|
]
|
|
|
|
|