AutoTrading/auto_trading/broker/backtest.py

import datetime

import pandas as pd  # type: ignore

from ..interfaces import DataBroker, CandlesProperties


class Backtest(DataBroker):
    """
    Backtest Broker
    """

    def __init__(self, f_name: str, start: int = 0, **kwargs) -> None:
        """Read the csv file and store it into a dataframe"""
        super().__init__()
        self.data = self._prepare_data(f_name, **kwargs)

        self.dates = self.data.index.levels[0]
        self._date0 = self.dates[0]
        self.step = self.calc_step()

        self._cursor = start
        self.start = start - 1

    def calc_step(self) -> datetime.datetime:
        return min(self.dates[i] - self.dates[i - 1] for i in range(1, 10))

    @property
    def cursor(self) -> datetime.datetime:
        return self.dates[self._cursor]

    @staticmethod
    def _prepare_data(f_name: str, **kwargs) -> pd.DataFrame:
        data = pd.read_csv(f_name, index_col=[0, 1], parse_dates=True, **kwargs)
        data.sort_index(inplace=True)
        data.fillna(method="ffill", inplace=True)
        return data

    @property
    def properties(self) -> CandlesProperties:
        """Return the properties of the candles for this broker."""
        return CandlesProperties(period=self.step)

    @property
    def current_change(self) -> pd.DataFrame:
        """Return the current change for each money."""
        return self.data.loc[(self.cursor,)]["close"].to_dict()

    def __iter__(self) -> "DataBroker":
        """Initialise the iterator."""
        self._cursor = self.start
        return self

    @property
    def __next__(self) -> pd.DataFrame:
        """Next values.

        Return the dataframe of all stock history for the strategy / indicators.

        Returns:
            DataFrame: Time-Stock valuated candlestick data.
                For each time and each stock give (high, low, open, close).
        """
        self._cursor += 1
        try:
            return self.data.loc[(self._date0,) : (self.cursor,)].copy  # type: ignore
        except IndexError:
            raise StopIteration