app.py

#!/usr/bin/env python3
"""
A股选股程序 - Web 服务
启动: python app.py
访问: http://localhost:5000
"""
import logging
import math
from flask import Flask, render_template, jsonify, request

from src.data_fetcher import DataFetcher
from src.screener import StockScreener
from strategies.builtin import BUILTIN_STRATEGIES

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)

app = Flask(__name__)

FIELDS_META = [
    {"field": "price", "label": "最新价", "unit": "元"},
    {"field": "change_pct", "label": "涨跌幅", "unit": "%"},
    {"field": "change", "label": "涨跌额", "unit": "元"},
    {"field": "volume", "label": "成交量", "unit": "手"},
    {"field": "amount", "label": "成交额", "unit": "元"},
    {"field": "amplitude", "label": "振幅", "unit": "%"},
    {"field": "high", "label": "最高价", "unit": "元"},
    {"field": "low", "label": "最低价", "unit": "元"},
    {"field": "open", "label": "今开价", "unit": "元"},
    {"field": "pre_close", "label": "昨收价", "unit": "元"},
    {"field": "volume_ratio", "label": "量比", "unit": ""},
    {"field": "turnover_rate", "label": "换手率", "unit": "%"},
    {"field": "pe_ttm", "label": "市盈率(TTM)", "unit": ""},
    {"field": "pb", "label": "市净率", "unit": ""},
    {"field": "total_mv", "label": "总市值", "unit": "元"},
    {"field": "circ_mv", "label": "流通市值", "unit": "元"},
    {"field": "change_pct_60d", "label": "60日涨跌幅", "unit": "%"},
    {"field": "change_pct_ytd", "label": "年初至今涨跌幅", "unit": "%"},
]

OPERATORS_META = [
    {"op": ">", "label": "大于 (>)"},
    {"op": ">=", "label": "大于等于 (>=)"},
    {"op": "<", "label": "小于 (<)"},
    {"op": "<=", "label": "小于等于 (<=)"},
    {"op": "==", "label": "等于 (==)"},
    {"op": "!=", "label": "不等于 (!=)"},
    {"op": "between", "label": "区间 (between)"},
]

STRATEGIES_META = [
    {"key": s_key, "label": s_desc}
    for s_key, (_, s_desc) in BUILTIN_STRATEGIES.items()
]

SORT_FIELDS = [
    {"field": "", "label": "不排序"},
    {"field": "price", "label": "最新价"},
    {"field": "change_pct", "label": "涨跌幅"},
    {"field": "volume_ratio", "label": "量比"},
    {"field": "turnover_rate", "label": "换手率"},
    {"field": "pe_ttm", "label": "市盈率(TTM)"},
    {"field": "pb", "label": "市净率"},
    {"field": "total_mv", "label": "总市值"},
    {"field": "circ_mv", "label": "流通市值"},
    {"field": "amount", "label": "成交额"},
    {"field": "change_pct_60d", "label": "60日涨跌幅"},
    {"field": "change_pct_ytd", "label": "年初至今涨跌幅"},
]


def _sanitize(val):
    """将 NaN / Inf 转为 None 以便 JSON 序列化"""
    if val is None:
        return None
    if isinstance(val, float) and (math.isnan(val) or math.isinf(val)):
        return None
    return val


@app.route("/")
def index():
    return render_template(
        "index.html",
        strategies=STRATEGIES_META,
        fields=FIELDS_META,
        operators=OPERATORS_META,
        sort_fields=SORT_FIELDS,
    )


@app.route("/api/meta")
def api_meta():
    """返回前端需要的元数据"""
    return jsonify({
        "strategies": STRATEGIES_META,
        "fields": FIELDS_META,
        "operators": OPERATORS_META,
        "sort_fields": SORT_FIELDS,
    })


@app.route("/api/screen", methods=["POST"])
def api_screen():
    """执行选股"""
    try:
        body = request.get_json(force=True)
        mode = body.get("mode", "strategy")  # "strategy" or "custom"
        strategy_name = body.get("strategy", "value_bluechip")
        conditions = body.get("conditions", [])
        exclude_st = body.get("exclude_st", True)
        exclude_suspended = body.get("exclude_suspended", True)
        sort_by = body.get("sort_by", "") or None
        ascending = body.get("ascending", False)
        top_n = body.get("top_n", 0) or None

        fetcher = DataFetcher()
        screener = StockScreener(fetcher=fetcher)

        if mode == "custom" and conditions:
            for c in conditions:
                value = c["value"]
                if c["operator"] == "between":
                    value = [float(value[0]), float(value[1])]
                else:
                    value = float(value)
                screener.add_condition(
                    field=c["field"],
                    operator=c["operator"],
                    value=value,
                    description=c.get("description", ""),
                )
        else:
            if strategy_name not in BUILTIN_STRATEGIES:
                return jsonify({"error": f"未知策略: {strategy_name}"}), 400
            strategy_func, _ = BUILTIN_STRATEGIES[strategy_name]
            strategy_func(screener)

        result = screener.screen(
            exclude_st=exclude_st,
            exclude_suspended=exclude_suspended,
            sort_by=sort_by,
            ascending=ascending,
            top_n=int(top_n) if top_n else None,
        )

        display_cols = [
            "code", "name", "price", "change_pct", "volume_ratio",
            "turnover_rate", "pe_ttm", "pb", "total_mv", "circ_mv",
            "amount", "amplitude", "change_pct_60d", "change_pct_ytd",
        ]
        cols = [c for c in display_cols if c in result.columns]
        rows = []
        for _, row in result[cols].iterrows():
            rows.append({c: _sanitize(row[c]) for c in cols})

        conditions_desc = [str(c) for c in screener.conditions]

        return jsonify({
            "total": len(rows),
            "conditions": conditions_desc,
            "data": rows,
        })
    except Exception as e:
        logging.exception("选股请求处理失败")
        return jsonify({"error": str(e)}), 500


@app.route("/api/strategy_conditions/<strategy_name>")
def api_strategy_conditions(strategy_name):
    """返回指定策略的筛选条件列表"""
    if strategy_name not in BUILTIN_STRATEGIES:
        return jsonify({"error": "未知策略"}), 404
    strategy_func, desc = BUILTIN_STRATEGIES[strategy_name]
    temp_screener = StockScreener.__new__(StockScreener)
    temp_screener.conditions = []
    temp_screener.fetcher = None
    temp_screener._data = None
    strategy_func(temp_screener)
    conditions = []
    for c in temp_screener.conditions:
        conditions.append({
            "field": c.field,
            "operator": c.operator,
            "value": c.value,
            "description": c.description,
        })
    return jsonify({"name": strategy_name, "description": desc, "conditions": conditions})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000, debug=False)