120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
"""
|
||
数据加载器 — 从SQLite加载K线数据为pandas DataFrame
|
||
"""
|
||
import sqlite3
|
||
import pandas as pd
|
||
from loguru import logger
|
||
from .config import DB_PATH, KLINE_PERIODS
|
||
|
||
|
||
def load_kline(period: int = 15, start_date: str = None, end_date: str = None) -> pd.DataFrame:
|
||
"""
|
||
加载指定周期的K线数据
|
||
:param period: K线周期（分钟），如 1, 3, 5, 15, 30, 60
|
||
:param start_date: 起始日期 'YYYY-MM-DD'（可选）
|
||
:param end_date: 结束日期 'YYYY-MM-DD'（可选）
|
||
:return: DataFrame，列: timestamp, open, high, low, close
|
||
"""
|
||
suffix = KLINE_PERIODS.get(period)
|
||
if suffix is None:
|
||
raise ValueError(f"不支持的周期: {period}，可选: {list(KLINE_PERIODS.keys())}")
|
||
|
||
table_name = f'bitmart_eth_{suffix}'
|
||
conn = sqlite3.connect(str(DB_PATH))
|
||
|
||
query = f"SELECT id as timestamp, open, high, low, close FROM {table_name} ORDER BY id"
|
||
df = pd.read_sql_query(query, conn)
|
||
conn.close()
|
||
|
||
if df.empty:
|
||
logger.warning(f"[{suffix}] 表中无数据")
|
||
return df
|
||
|
||
# id 是毫秒时间戳，转为 datetime 索引
|
||
df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms')
|
||
df.set_index('datetime', inplace=True)
|
||
|
||
# 按日期过滤
|
||
if start_date:
|
||
df = df[df.index >= start_date]
|
||
if end_date:
|
||
df = df[df.index <= end_date]
|
||
|
||
logger.info(f"[{suffix}] 加载 {len(df)} 条K线 | {df.index[0]} ~ {df.index[-1]}")
|
||
return df
|
||
|
||
|
||
def load_multi_period(periods: list = None, start_date: str = None, end_date: str = None) -> dict:
|
||
"""
|
||
加载多个周期的K线数据
|
||
:param periods: 周期列表，如 [5, 15, 60]，默认全部
|
||
:param start_date: 起始日期
|
||
:param end_date: 结束日期
|
||
:return: {period: DataFrame} 字典
|
||
"""
|
||
if periods is None:
|
||
periods = list(KLINE_PERIODS.keys())
|
||
|
||
result = {}
|
||
for p in periods:
|
||
try:
|
||
df = load_kline(p, start_date, end_date)
|
||
if not df.empty:
|
||
result[p] = df
|
||
except Exception as e:
|
||
logger.error(f"加载 {p}分钟 K线失败: {e}")
|
||
|
||
return result
|
||
|
||
|
||
def load_trades(start_date: str = None, end_date: str = None, limit: int = None) -> pd.DataFrame:
|
||
"""
|
||
加载原始成交记录
|
||
:return: DataFrame，列: id, timestamp, price, volume, side
|
||
"""
|
||
conn = sqlite3.connect(str(DB_PATH))
|
||
query = "SELECT id, timestamp, price, volume, side FROM bitmart_eth_trades ORDER BY timestamp"
|
||
df = pd.read_sql_query(query, conn)
|
||
conn.close()
|
||
|
||
if df.empty:
|
||
logger.warning("成交记录表中无数据")
|
||
return df
|
||
|
||
df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms')
|
||
df.set_index('datetime', inplace=True)
|
||
|
||
if start_date:
|
||
df = df[df.index >= start_date]
|
||
if end_date:
|
||
df = df[df.index <= end_date]
|
||
if limit:
|
||
df = df.head(limit)
|
||
|
||
logger.info(f"加载 {len(df)} 条成交记录")
|
||
return df
|
||
|
||
|
||
def get_available_tables() -> list:
|
||
"""列出数据库中所有可用的表"""
|
||
conn = sqlite3.connect(str(DB_PATH))
|
||
cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
|
||
tables = [row[0] for row in cursor.fetchall()]
|
||
conn.close()
|
||
return tables
|
||
|
||
|
||
def get_table_stats() -> dict:
|
||
"""获取各表的数据统计"""
|
||
conn = sqlite3.connect(str(DB_PATH))
|
||
tables = get_available_tables()
|
||
stats = {}
|
||
for table in tables:
|
||
try:
|
||
count = pd.read_sql_query(f"SELECT COUNT(*) as cnt FROM {table}", conn).iloc[0]['cnt']
|
||
stats[table] = count
|
||
except Exception:
|
||
stats[table] = 0
|
||
conn.close()
|
||
return stats
|