"""generate_demand_agent 分类树热度维度共享常量与工具函数。""" from __future__ import annotations from collections.abc import Callable from decimal import Decimal from supply_infra.db.models.category_tree_weight import CategoryTreeWeight from supply_infra.db.models.global_tree_category import GlobalTreeCategory DIM_KEYS: tuple[str, ...] = ( "ext_pop", "plat_sust_pop", "plat_ly_pop", "recent_pop", ) DIM_LABEL: dict[str, str] = { "ext_pop": "外部热度", "plat_sust_pop": "平台持续热度", "plat_ly_pop": "平台去年同期热度", "recent_pop": "近期热度", } # 节点下存在有数据的叶子节点时的标记 HAS_DATA_LEAF_MARK = "+" def normalize_biz_dt(biz_dt: str | None) -> tuple[str | None, str | None]: """校验并规范化 biz_dt(YYYYMMDD);空值表示使用最新业务日。""" if biz_dt is None: return None, None text = str(biz_dt).strip() if not text: return None, None if len(text) != 8 or not text.isdigit(): return None, f"biz_dt 格式无效,应为 YYYYMMDD: {biz_dt!r}" return text, None def resolve_biz_dt( biz_dt: str | None, *, get_latest: Callable[[], str | None], has_data: Callable[[str], bool] | None = None, table_label: str = "热度数据", ) -> tuple[str | None, str | None]: """解析 biz_dt:未传则用最新业务日;传入则校验格式与数据是否存在。""" normalized, err = normalize_biz_dt(biz_dt) if err: return None, err resolved = normalized or get_latest() if not resolved: return None, f"暂无 {table_label}(请指定 biz_dt 或先写入数据)" if normalized and has_data is not None and not has_data(normalized): return None, f"biz_dt={normalized} 无 {table_label}" return resolved, None def get_latest_common_biz_dt(session) -> str | None: """返回 category_tree_weight 与 demand_popularity_stats 均有数据的最新 biz_dt。""" from sqlalchemy import func, select from supply_infra.db.models.demand_popularity_stats import DemandPopularityStats stats_dates = select(DemandPopularityStats.biz_dt).distinct() stmt = select(func.max(CategoryTreeWeight.biz_dt)).where( CategoryTreeWeight.biz_dt.in_(stats_dates) ) return session.scalar(stmt) 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 format_score(avg: Decimal | float | int | None) -> str: if avg is None: return "—" value = float(avg) if value >= 100: return f"{value:.0f}" if value >= 1: return f"{value:.1f}" return f"{value:.2f}" def dim_score( weight: CategoryTreeWeight | None, dim: str, ) -> tuple[float | None, int]: """返回 (avg, count);count>0 即有维度数据。""" if weight is None: return None, 0 count = int(getattr(weight, f"{dim}_count", 0) or 0) if count <= 0: return None, 0 avg = getattr(weight, f"{dim}_avg", None) return (float(avg) if avg is not None else None), count 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 build_score_by_id( weights: list[CategoryTreeWeight], dim: str, ) -> dict[int, tuple[float | None, int]]: return {int(row.category_id): dim_score(row, dim) for row in weights} def collect_descendant_leaves( root_id: int, children_map: dict[int | None, list[GlobalTreeCategory]], ) -> list[int]: """收集 root 子树内的叶子节点 id(若 root 本身无子节点则含 root)。""" leaves: list[int] = [] stack = [root_id] seen: set[int] = set() while stack: cid = stack.pop() if cid in seen: continue seen.add(cid) kids = children_map.get(cid, []) if not kids: leaves.append(cid) else: stack.extend(int(c.id) for c in kids) leaves.sort() return leaves def collect_descendant_ids( root_id: int, children_map: dict[int | None, list[GlobalTreeCategory]], ) -> list[int]: """收集 root 及其全部子孙节点 id。""" ids: list[int] = [] stack = [root_id] seen: set[int] = set() while stack: cid = stack.pop() if cid in seen: continue seen.add(cid) ids.append(cid) stack.extend(int(c.id) for c in children_map.get(cid, [])) ids.sort() return ids def build_category_path( category_id: int, by_id: dict[int, GlobalTreeCategory], ) -> str | None: """从根到指定节点的名称路径,如「美妆护肤 > 护肤 > 防晒」。""" cat = by_id.get(category_id) if cat is None: return None names: list[str] = [] current: GlobalTreeCategory | None = cat seen: set[int] = set() while current is not None: cid = int(current.id) if cid in seen: break seen.add(cid) names.append(current.name or str(cid)) parent_key = normalize_parent_id(current.parent_id) current = by_id.get(parent_key) if parent_key is not None else None names.reverse() return " > ".join(names)