124 lines
4.2 KiB
Python
124 lines
4.2 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import logging
|
||
from datetime import datetime, timedelta
|
||
from sqlalchemy import create_engine, text
|
||
import pandas as pd
|
||
|
||
from src.valuation_analysis.config import DB_URL
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class IndexAnalyzer:
|
||
"""指数数据分析工具类"""
|
||
|
||
def __init__(self, db_url=None):
|
||
# 初始化数据库连接
|
||
self.db_url = db_url or DB_URL
|
||
self.engine = create_engine(self.db_url)
|
||
|
||
def get_indices_list(self):
|
||
"""
|
||
获取可用指数列表
|
||
|
||
Returns:
|
||
list: 包含指数信息的列表 [{"id": id, "name": name, "code": code}, ...]
|
||
"""
|
||
try:
|
||
with self.engine.connect() as conn:
|
||
query = text("""
|
||
SELECT id, gp_name as name, gp_code as code
|
||
FROM gp_code_zs
|
||
ORDER BY gp_name
|
||
""")
|
||
result = conn.execute(query).fetchall()
|
||
|
||
indices = []
|
||
for row in result:
|
||
indices.append({
|
||
"id": row[0],
|
||
"name": row[1],
|
||
"code": row[2]
|
||
})
|
||
|
||
logger.info(f"获取到 {len(indices)} 个指数")
|
||
return indices
|
||
except Exception as e:
|
||
logger.error(f"获取指数列表失败: {str(e)}")
|
||
return []
|
||
|
||
def get_index_data(self, index_code, start_date=None, end_date=None):
|
||
"""
|
||
获取指数历史数据
|
||
|
||
Args:
|
||
index_code: 指数代码
|
||
start_date: 开始日期 (可选，默认为1年前)
|
||
end_date: 结束日期 (可选，默认为今天)
|
||
|
||
Returns:
|
||
dict: 包含指数数据的字典 {"code": code, "dates": [...], "values": [...]}
|
||
"""
|
||
try:
|
||
# 处理日期参数
|
||
if end_date is None:
|
||
end_date = datetime.now().strftime('%Y-%m-%d')
|
||
|
||
if start_date is None:
|
||
start_date = (datetime.now() - timedelta(days=365)).strftime('%Y-%m-%d')
|
||
|
||
with self.engine.connect() as conn:
|
||
query = text("""
|
||
SELECT timestamp, close
|
||
FROM gp_day_data
|
||
WHERE symbol = :symbol
|
||
AND timestamp BETWEEN :start_date AND :end_date
|
||
ORDER BY timestamp
|
||
""")
|
||
|
||
result = conn.execute(query, {
|
||
"symbol": index_code,
|
||
"start_date": start_date,
|
||
"end_date": end_date
|
||
}).fetchall()
|
||
|
||
dates = []
|
||
values = []
|
||
|
||
for row in result:
|
||
dates.append(row[0].strftime('%Y-%m-%d'))
|
||
# close可能是字符串类型，转换为浮点数
|
||
values.append(float(row[1]) if row[1] else None)
|
||
|
||
logger.info(f"获取指数 {index_code} 数据: {len(dates)} 条记录")
|
||
return {
|
||
"code": index_code,
|
||
"dates": dates,
|
||
"values": values
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"获取指数 {index_code} 数据失败: {str(e)}")
|
||
return {
|
||
"code": index_code,
|
||
"dates": [],
|
||
"values": []
|
||
}
|
||
|
||
# 测试代码
|
||
if __name__ == "__main__":
|
||
analyzer = IndexAnalyzer()
|
||
|
||
# 测试获取指数列表
|
||
indices = analyzer.get_indices_list()
|
||
print(f"指数列表: {indices[:5]}...")
|
||
|
||
# 测试获取指数数据
|
||
if indices:
|
||
# 测试第一个指数的数据
|
||
first_index = indices[0]
|
||
index_data = analyzer.get_index_data(first_index['code'])
|
||
print(f"指数 {first_index['name']} 数据:")
|
||
print(f"日期数量: {len(index_data['dates'])}")
|
||
if index_data['dates']:
|
||
print(f"第一个日期: {index_data['dates'][0]}, 值: {index_data['values'][0]}") |