70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
import os
|
|
import threading
|
|
from typing import List, Optional
|
|
|
|
import pandas as pd
|
|
import pyarrow.dataset as ds
|
|
|
|
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def ensure_storage(base_dir: str):
|
|
os.makedirs(base_dir, exist_ok=True)
|
|
|
|
|
|
def _path(base_dir: str, symbol: str, timeframe: str) -> str:
|
|
safe_symbol = symbol.replace("/", "_").replace(":", "_")
|
|
d = os.path.join(base_dir, timeframe)
|
|
os.makedirs(d, exist_ok=True)
|
|
return os.path.join(d, f"{safe_symbol}.parquet")
|
|
|
|
|
|
def read_candles(base_dir: str, symbol: str, timeframe: str, start: Optional[int], end: Optional[int]) -> pd.DataFrame:
|
|
p = _path(base_dir, symbol, timeframe)
|
|
if not os.path.exists(p):
|
|
return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) # empty
|
|
df = pd.read_parquet(p)
|
|
if start is not None:
|
|
df = df[df["timestamp"] >= int(start)]
|
|
if end is not None:
|
|
df = df[df["timestamp"] <= int(end)]
|
|
df = df.sort_values("timestamp")
|
|
return df
|
|
|
|
|
|
def upsert_candles(base_dir: str, symbol: str, timeframe: str, candles: List[List[float]]):
|
|
p = _path(base_dir, symbol, timeframe)
|
|
new_df = pd.DataFrame(candles, columns=["timestamp", "open", "high", "low", "close", "volume"])
|
|
with _lock:
|
|
if os.path.exists(p):
|
|
old = pd.read_parquet(p)
|
|
merged = pd.concat([old, new_df], ignore_index=True)
|
|
merged = merged.drop_duplicates(subset=["timestamp"], keep="last").sort_values("timestamp")
|
|
else:
|
|
merged = new_df.sort_values("timestamp")
|
|
merged.to_parquet(p, index=False)
|
|
|
|
|
|
def get_last_timestamp(base_dir: str, symbol: str, timeframe: str) -> Optional[int]:
|
|
p = _path(base_dir, symbol, timeframe)
|
|
if not os.path.exists(p):
|
|
return None
|
|
df = pd.read_parquet(p)
|
|
if df.empty:
|
|
return None
|
|
return int(df["timestamp"].iloc[-1])
|
|
|
|
|
|
def read_candle_exact(base_dir: str, symbol: str, timeframe: str, timestamp: int) -> pd.DataFrame:
|
|
p = _path(base_dir, symbol, timeframe)
|
|
if not os.path.exists(p):
|
|
return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"])
|
|
dataset = ds.dataset(p, format="parquet")
|
|
table = dataset.to_table(filter=ds.field("timestamp") == int(timestamp))
|
|
if table.num_rows == 0:
|
|
return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"])
|
|
return table.to_pandas()
|
|
|
|
|