54 lines
1.2 KiB
Python
54 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import pandas as pd
|
|
import sqlite3
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class KlineSource:
|
|
db_path: Path
|
|
table_name: str
|
|
|
|
|
|
def _to_ms(dt: datetime) -> int:
|
|
if dt.tzinfo is None:
|
|
dt = dt.replace(tzinfo=timezone.utc)
|
|
return int(dt.timestamp() * 1000)
|
|
|
|
|
|
def load_klines(
|
|
source: KlineSource,
|
|
start: datetime,
|
|
end: datetime,
|
|
) -> pd.DataFrame:
|
|
start_ms = _to_ms(start)
|
|
end_ms = _to_ms(end)
|
|
|
|
con = sqlite3.connect(str(source.db_path))
|
|
try:
|
|
df = pd.read_sql_query(
|
|
f"SELECT id, open, high, low, close FROM {source.table_name} WHERE id >= ? AND id <= ? ORDER BY id ASC",
|
|
con,
|
|
params=(start_ms, end_ms),
|
|
)
|
|
finally:
|
|
con.close()
|
|
|
|
if df.empty:
|
|
return df
|
|
|
|
df["timestamp_ms"] = df["id"].astype("int64")
|
|
df["dt"] = pd.to_datetime(df["timestamp_ms"], unit="ms", utc=True)
|
|
df = df.drop(columns=["id"]).set_index("dt")
|
|
|
|
for c in ("open", "high", "low", "close"):
|
|
df[c] = pd.to_numeric(df[c], errors="coerce")
|
|
|
|
df = df.dropna(subset=["open", "high", "low", "close"])
|
|
return df
|