fa-agent/finance/formula_columns.py

import math
from decimal import Decimal, localcontext

from .analysis import METRICS

RESERVED_DATA_HEADERS = {
    "period",
    "customertoken",
    "revenue",
    "cost",
    "grossinput",
    "marginrate",
    "grosssource",
    "listamount",
    "selfoperated",
    "customerpresent",
    "rowcount",
    "grosscalc",
}


def available_metrics(ds):
    roles = {column["role"] for column in ds.mapping["columns"]}
    dimensions = {column.get("key") for column in ds.mapping["columns"] if column["role"] == "dimension"}
    products = bool(ds.mapping.get("self_operated_products")) and "product" in dimensions
    gross = "gross" in roles or (bool(ds.mapping.get("margin_rates")) and "product" in dimensions)
    customer, list_amount, cost = "customer" in roles, "list_amount" in roles, "cost" in roles
    available = {"revenue"}
    if cost:
        available.add("cost")
    if gross:
        available.update(("gross", "margin"))
    if customer:
        available.update(("customers", "basket"))
    if list_amount:
        available.add("discount_rate")
    if customer and list_amount:
        available.add("list_basket")
    if cost and (ds.mapping.get("roi_formula") == "profit_over_cost" or gross):
        available.add("roi")
    if products:
        available.add("self_revenue")
        if customer:
            available.add("upgrade_rate")
    return {key: value for key, value in METRICS.items() if key in available}


OPERATORS = {"add": "+", "subtract": "-", "multiply": "*", "divide": "/"}


def validate_column(value):
    if not isinstance(value, dict) or set(value) != {"name", "meaning", "scope", "missing", "expression"}:
        raise ValueError("自定义列定义字段无效")
    for key, limit in (("name", 80), ("meaning", 500), ("scope", 200), ("missing", 300)):
        text = value[key]
        if not isinstance(text, str) or not text.strip() or len(text) > limit:
            raise ValueError(f"自定义列{key}说明不完整")
    if value["missing"] != "结果不可用（#N/A）；零分母不可用":
        raise ValueError("暂不支持自定义缺失值策略；请选择结果不可用（#N/A）及零分母不可用")
    _validate_expression(value["expression"])
    return value


def _validate_expression(node, depth=0):
    if depth > 8 or not isinstance(node, dict):
        raise ValueError("自定义计算公式结构无效")
    if set(node) == {"metric"}:
        if not isinstance(node["metric"], str) or node["metric"] not in METRICS:
            raise ValueError("公式引用了不支持的指标")
        return
    if set(node) == {"number"}:
        if (
            type(node["number"]) not in (int, float)
            or not math.isfinite(node["number"])
            or abs(node["number"]) > 1e12
        ):
            raise ValueError("公式常数无效")
        return
    if (
        set(node) == {"operator", "left", "right"}
        and isinstance(node["operator"], str)
        and node["operator"] in OPERATORS
    ):
        _validate_expression(node["left"], depth + 1)
        _validate_expression(node["right"], depth + 1)
        return
    raise ValueError("公式包含不支持的字段或运算")


def metrics_used(node):
    if "metric" in node:
        return {node["metric"]}
    if "number" in node:
        return set()
    return metrics_used(node["left"]) | metrics_used(node["right"])


def ensure_available_metrics(ds, column):
    unavailable = metrics_used(column["expression"]) - available_metrics(ds).keys()
    if unavailable:
        names = "、".join(METRICS[key] for key in sorted(unavailable))
        raise ValueError(f"当前数据集未映射计算“{names}”所需字段，不能生成此自定义列")


def to_excel(node, references):
    if "metric" in node:
        return references[node["metric"]]
    if "number" in node:
        return str(node["number"])
    return f"({to_excel(node['left'], references)}{OPERATORS[node['operator']]}{to_excel(node['right'], references)})"


def evaluate(node, references):
    """Evaluate a validated expression with Decimal values and #N/A semantics."""
    with localcontext() as decimal_context:
        decimal_context.prec = 64
        return _evaluate(node, references)


def _evaluate(node, references):
    if "metric" in node:
        return references.get(node["metric"])
    if "number" in node:
        return Decimal(str(node["number"]))
    left, right = _evaluate(node["left"], references), _evaluate(node["right"], references)
    if left is None or right is None:
        return None
    if node["operator"] == "add":
        return left + right
    if node["operator"] == "subtract":
        return left - right
    if node["operator"] == "multiply":
        return left * right
    return None if right == 0 else left / right