_common/data/repository/KlineRepository.py - dev - fgtmenow/tb

_common/data/repository/KlineRepository.py
179 строк6 KB
Azarev Artem
market structure, choch, confirm, bos
02 апр 2025, 14:24
02 апр 2025, 14:2487ab4da
import datetime as dt
import os
from concurrent.futures import ThreadPoolExecutor

import pandas as pd
from pandas import DataFrame
from pybit.unified_trading import HTTP

from _common.api.GetKline import get_kline
from _common.domain.model.BacktestConfig import BacktestConfig
from _common.domain.model.Instrument import Instrument
from _common.domain.util.CacheUtil import get_cache_file_path
from _common.domain.util.Logger import print_message

CACHE_DIR = "__cache"


def load_data(
        session: HTTP,
        symbol: str,
        timeframe: str,
        period: BacktestConfig,
        cache: bool,
        save_to_cache: bool = True
) -> DataFrame:
    cache_file_name = get_cache_file_path(symbol, "frames", period.config_name)
    if cache and os.path.exists(cache_file_name):
        df = pd.read_csv(cache_file_name)
        df = df.drop_duplicates(subset=['timestamp'], keep='last')
        # Convert timestamp to datetime index
        df.index = pd.to_datetime(df['timestamp'], unit='ms')
        df = df.sort_index()
    else:
        df = __download_data(
            session=session,
            symbol=symbol,
            timeframe=timeframe,
            start=period.start,
            end=None
        )
        if save_to_cache:
            df.to_csv(get_cache_file_path(symbol, "frames", period.config_name))
    # add DataFrame to result
    df["symbol"] = symbol
    return df


def __download_data(session: HTTP,
                    symbol: str,
                    timeframe: str,
                    start: int,
                    end: int | None,
                    ) -> pd.DataFrame:
    df = pd.DataFrame()

    chunk_start = start
    while True:
        latest = get_kline(session, symbol, chunk_start, timeframe, 1000)
        if not isinstance(latest, pd.DataFrame):
            break
        chunk_start = int(get_last_timestamp(latest))
        df = pd.concat([df, latest])

        if end is not None and chunk_start >= end:
            break

        start_date = dt.datetime.fromtimestamp(chunk_start / 1000).strftime('%d.%m.%Y %H:%M')
        print_message(f"Collect {symbol} from {start_date}, collected {len(df.index)}")
        if len(latest) == 1:
            break
    df.close = df.close.apply(lambda x: float(x))
    df.high = df.high.apply(lambda x: float(x))
    df.low = df.low.apply(lambda x: float(x))
    df.volume = df.volume.apply(lambda x: float(x))
    df.open = df.open.apply(lambda x: float(x))
    df.timestamp = df.timestamp.apply(lambda x: int(x))
    if end is not None:
        df = df[df['timestamp'] <= end]
    df = df.drop_duplicates(subset=['timestamp'], keep='last')
    # Convert timestamp to datetime index
    df.index = pd.to_datetime(df['timestamp'], unit='ms')
    df.index.name = "index"
    df = df.sort_index()
    return df


def get_cached_kline_data_frames(
        session: HTTP,
        timeframe: str,
        instruments: dict[str, Instrument],
        period: BacktestConfig
) -> dict[str, DataFrame]:
    symbols = list(instruments.keys())
    result: dict[str, DataFrame] = {}
    for symbol in symbols:
        result[symbol] = load_data(session, symbol, timeframe, period, True)

    # Получить минимальный общий начальный и максимальный общий конечный timestamp
    start_timestamp = max(df['timestamp'].min() for df in result.values())
    end_timestamp = min(df['timestamp'].max() for df in result.values())

    # Преобразовать int timestamp в datetime
    start_timestamp_dt = pd.to_datetime(start_timestamp / 1000, unit='s')
    end_timestamp_dt = pd.to_datetime(end_timestamp / 1000, unit='s')

    # Преобразовать timestamp в формат %d.%m.%Y %H:%M
    start_timestamp_formatted = start_timestamp_dt.strftime('%d.%m.%Y %H:%M')
    end_timestamp_formatted = end_timestamp_dt.strftime('%d.%m.%Y %H:%M')

    df_trimmed: dict[str, DataFrame] = {}
    for symbol, df in result.items():
        trimmed_df = df[(df['timestamp'] >= start_timestamp) & (df['timestamp'] <= end_timestamp)]
        df_trimmed[symbol] = trimmed_df
    return df_trimmed


def get_kline_data_frame(
        session: HTTP,
        timeframe: str,
        instruments: dict[str, Instrument],
        period: BacktestConfig,
        cache: bool
):
    symbols = list(instruments.keys())
    with ThreadPoolExecutor(40) as executor:
        # Метод map возвращает результаты в том же порядке, что и параметры

        result = list(
            executor.map(load_data,
                         [session] * len(symbols),
                         symbols,
                         [timeframe] * len(symbols),
                         [period] * len(symbols),
                         [cache] * len(symbols)
                         )
        )

    # Получить минимальный общий начальный и максимальный общий конечный timestamp
    start_timestamp = max(df['timestamp'].min() for df in result)
    end_timestamp = min(df['timestamp'].max() for df in result)

    # Преобразовать int timestamp в datetime
    start_timestamp_dt = pd.to_datetime(start_timestamp / 1000, unit='s')
    end_timestamp_dt = pd.to_datetime(end_timestamp / 1000, unit='s')

    # Преобразовать timestamp в формат %d.%m.%Y %H:%M
    start_timestamp_formatted = start_timestamp_dt.strftime('%d.%m.%Y %H:%M')
    end_timestamp_formatted = end_timestamp_dt.strftime('%d.%m.%Y %H:%M')

    df_trimmed = []
    for df in result:
        trimmed_df = df[(df['timestamp'] >= start_timestamp) & (df['timestamp'] <= end_timestamp)]
        df_trimmed.append(trimmed_df)
    print_message(f"DataFrames trimmed from {start_timestamp_formatted} to {end_timestamp_formatted}")
    return df_trimmed


def get_last_timestamp(df):
    return int(df.timestamp.values[0])


def format_data(response):
    data = response.get('list', None)
    if not data:
        return
    df = pd.DataFrame(data,
                      columns=[
                          'timestamp',
                          'open',
                          'high',
                          'low',
                          'close',
                          'volume',
                          'turnover'
                      ],
                      )
    # Set datetime index directly from timestamp
    df.index = pd.to_datetime(df['timestamp'], unit='ms')
    return df