"""Build nested category tree from global_tree_category (+ optional weights).""" from __future__ import annotations from typing import Any from supply_infra.db.models.category_tree_weight import CategoryTreeWeight from supply_infra.db.models.global_tree_category import GlobalTreeCategory from supply_infra.db.repositories.category_tree_weight_repo import ( CategoryTreeWeightRepository, ) from supply_infra.db.repositories.global_tree_category_repo import ( GlobalTreeCategoryRepository, ) from supply_infra.db.session import get_session DIM_KEYS: tuple[str, ...] = ( "ext_pop", "plat_sust_pop", "plat_ly_pop", "recent_pop", "real_rov_7d", "real_vov_7d", ) TOTAL_SCORE_KEY = "total_score" TOTAL_SCORE_DIM_KEYS: tuple[str, ...] = ( "ext_pop", "plat_sust_pop", "plat_ly_pop", "recent_pop", ) DIM_META: list[dict[str, str]] = [ {"key": TOTAL_SCORE_KEY, "label": "全局热度"}, {"key": "ext_pop", "label": "外部热度"}, {"key": "plat_sust_pop", "label": "平台持续热度"}, {"key": "plat_ly_pop", "label": "去年同期热度"}, {"key": "recent_pop", "label": "近期热度"}, {"key": "real_rov_7d", "label": "真实ROV(7日)"}, {"key": "real_vov_7d", "label": "真实VOV(7日)"}, ] def _normalize_parent_id(parent_id: int | None) -> int | None: if parent_id is None or parent_id == 0: return None return parent_id def _weights_payload(row: CategoryTreeWeight | None) -> dict[str, float | None]: if row is None: return {dim: None for dim in (TOTAL_SCORE_KEY, *DIM_KEYS)} return { TOTAL_SCORE_KEY: float(row.total_score) if row.total_score is not None else None, **{ dim: ( float(getattr(row, f"{dim}_avg")) if getattr(row, f"{dim}_avg", None) is not None else None ) for dim in DIM_KEYS }, } def _counts_payload(row: CategoryTreeWeight | None) -> dict[str, int]: if row is None: return {dim: 0 for dim in (TOTAL_SCORE_KEY, *DIM_KEYS)} return { TOTAL_SCORE_KEY: sum( int(getattr(row, f"{dim}_count", 0) or 0) > 0 for dim in TOTAL_SCORE_DIM_KEYS ), **{ dim: int(getattr(row, f"{dim}_count", 0) or 0) for dim in DIM_KEYS }, } def _build_children_map( categories: list[GlobalTreeCategory], ) -> dict[int | None, list[GlobalTreeCategory]]: children_map: dict[int | None, list[GlobalTreeCategory]] = {} for category in categories: parent_key = _normalize_parent_id(category.parent_id) children_map.setdefault(parent_key, []).append(category) for children in children_map.values(): children.sort(key=lambda c: (c.level or 0, c.id)) return children_map def _to_node( category: GlobalTreeCategory, children_map: dict[int | None, list[GlobalTreeCategory]], weight_by_id: dict[int, CategoryTreeWeight], ) -> dict[str, Any]: children = [ _to_node(child, children_map, weight_by_id) for child in children_map.get(category.id, []) ] weight_row = weight_by_id.get(int(category.id)) return { "id": category.id, "name": category.name, "level": category.level, "description": category.description, "weights": _weights_payload(weight_row), "counts": _counts_payload(weight_row), "hung_word_count": int(weight_row.hung_word_count) if weight_row else 0, "children": children, } def build_category_tree(biz_dt: str | None = None) -> dict[str, Any]: """Load active categories with per-dim avg; nest as roots. Returns ``{biz_dt, dims, nodes}``. When no weight rows exist, ``biz_dt`` is ``None`` and each node still carries null ``weights`` / zero ``counts``. """ with get_session() as session: categories = GlobalTreeCategoryRepository(session).list_active_categories() weight_repo = CategoryTreeWeightRepository(session) resolved_dt = biz_dt or weight_repo.get_latest_biz_dt() weight_rows: list[CategoryTreeWeight] = [] if resolved_dt: weight_rows = weight_repo.list_by_biz_dt(resolved_dt) weight_by_id = {int(r.category_id): r for r in weight_rows} children_map = _build_children_map(categories) nodes = [ _to_node(root, children_map, weight_by_id) for root in children_map.get(None, []) ] return { "biz_dt": resolved_dt, "dims": DIM_META, "nodes": nodes, }