diff --git a/core/datasource/flask_server.py b/core/datasource/flask_server.py index d9874f7..e1e8fba 100644 --- a/core/datasource/flask_server.py +++ b/core/datasource/flask_server.py @@ -232,7 +232,8 @@ def dataframe_to_json(df: pd.DataFrame) -> Dict: # 转换为字典列表 records = df_reset.to_dict(orient='records') - return { + # 构建返回结果 + result = { "data": records, "count": len(records), "columns": list(df_reset.columns), @@ -241,6 +242,12 @@ def dataframe_to_json(df: pd.DataFrame) -> Dict: "end": df.index.max().strftime('%Y-%m-%d') if hasattr(df.index.max(), 'strftime') else str(df.index.max()), } if len(df) > 0 else None } + + # 添加股票信息(如果存在) + if hasattr(df, 'attrs') and df.attrs.get('info'): + result['info'] = df.attrs['info'] + + return result def validate_date(date_str: str) -> bool: diff --git a/core/datasource/universal_fetcher.py b/core/datasource/universal_fetcher.py index 0da1fe7..1fc9eb7 100644 --- a/core/datasource/universal_fetcher.py +++ b/core/datasource/universal_fetcher.py @@ -362,6 +362,20 @@ class UniversalDataFetcher: try: ticker = yf.Ticker(yf_code) + + # 获取公司信息(仅对股票) + info = {} + if asset_type in ['us_stock', 'hk_stock']: + try: + stock_info = ticker.info + info = { + 'sector': stock_info.get('sector'), + 'industry': stock_info.get('industry'), + 'market_cap': stock_info.get('marketCap'), + } + except Exception: + pass + # end_date 需要加一天(yfinance 的 end 是排他的) end_date_obj = pd.Timestamp(end_date) + timedelta(days=1) data = ticker.history( @@ -388,6 +402,10 @@ class UniversalDataFetcher: data = data[available] data['code'] = code + # 添加公司信息到 DataFrame 的 attrs(属性) + if info: + data.attrs['info'] = info + return data except Exception as e: