用Fastapi写个ETF计算应用

代码如下

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import pandas as pd
import akshare as ak
from datetime import datetime

# ETF代码与名称的映射
etf_names = {
    '510300': '300ETF',
    '159915': '创业板',
    '513050': '中概互联网ETF',
    '159941': '纳指ETF',
    '518880': '黄金ETF',
    '511090': '30年国债ETF'
}

# 获取ETF数据的函数
def fetch_etf_data(symbol="510300"):
    df = ak.fund_etf_hist_em(symbol=symbol, period="daily", adjust='qfq')
    df['日期'] = pd.to_datetime(df['日期'])
    df.set_index('日期', inplace=True)
    df.rename(columns={
        '开盘': 'Open',
        '最高': 'High',
        '最低': 'Low',
        '收盘': 'Close',
        '成交量': 'Volume'
    }, inplace=True)
    return df

# 计算动量和均线
def calculate_momentum_and_ma(df, momentum_period=20, ma_period=28):
    df['Momentum'] = df['Close'] / df['Close'].shift(momentum_period) - 1
    df['MA'] = df['Close'].rolling(window=ma_period).mean()
    return df

# 筛选符合条件的ETF
def select_etfs(etf_list, momentum_period=20, ma_period=28):
    etf_data = {}
    for symbol in etf_list:
        df = fetch_etf_data(symbol)
        df = calculate_momentum_and_ma(df, momentum_period, ma_period)
        etf_data[symbol] = df

    latest_data = {symbol: df.iloc[-1] for symbol, df in etf_data.items()}
    all_etfs = []
    for symbol, data in latest_data.items():
        above_ma = data['Close'] > data['MA']
        all_etfs.append((symbol, etf_names[symbol], data['Momentum'], above_ma))

    all_etfs.sort(key=lambda x: x[2], reverse=True)
    selected_etfs = [(etf[0], etf[1], etf[2]) for etf in all_etfs if etf[3]][:2]

    return selected_etfs, all_etfs

# FastAPI应用
app = FastAPI()

# 启用CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源访问
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ETFResponse(BaseModel):
    date: str
    recommended_etfs: list
    all_etfs: list

@app.get("/analyze_etfs", response_model=ETFResponse)
def analyze_etfs():
    try:
        etf_list = ['510300', '159915', '513050', '159941', '518880', '511090']
        selected_etfs, all_etfs = select_etfs(etf_list)

        formatted_selected_etfs = [
            {"symbol": symbol, "name": name, "momentum": float(momentum)}
            for symbol, name, momentum in selected_etfs
        ]

        formatted_all_etfs = [
            {
                "symbol": symbol,
                "name": name,
                "momentum": float(momentum),
                "above_ma": bool(above_ma)  # 显式转换为Python布尔类型
            }
            for symbol, name, momentum, above_ma in all_etfs
        ]

        response = {
            "date": datetime.now().strftime("%Y-%m-%d"),
            "recommended_etfs": formatted_selected_etfs,
            "all_etfs": formatted_all_etfs
        }

        return response

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)