progress

src/strategies.py

@@ -1,10 +1,9 @@
 from abc import ABC, abstractmethod
 from enum import Enum
-import sys
 
-from typing import Dict, Optional
+from typing import Dict, Optional, cast
 
-import pandas as pd  # type: ignore
+import pandas as pd  # type: ignore[import]
 
 from tools.trading_pair import TradingPair
 from results import BacktestResult
@@ -24,6 +23,10 @@ class PairsTradingStrategy(ABC):
     @abstractmethod
     def run_pair(self, config: Dict, pair: TradingPair, bt_result: BacktestResult) -> Optional[pd.DataFrame]:
         ...
+        
+    @abstractmethod
+    def reset(self):
+        ...
 
 class StaticFitStrategy(PairsTradingStrategy):
 
@@ -197,8 +200,11 @@ class StaticFitStrategy(PairsTradingStrategy):
         # Add tuples to data frame
         return pd.DataFrame(
             trd_signal_tuples,
-            columns=self.TRADES_COLUMNS,
+            columns=self.TRADES_COLUMNS,  # type: ignore
         )
+    
+    def reset(self):
+        pass
 
 class PairState(Enum):
     INITIAL = 1
@@ -214,7 +220,7 @@ class SlidingFitStrategy(PairsTradingStrategy):
         print(f"***{pair}*** STARTING....")
 
         pair.user_data_['state'] = PairState.INITIAL
-        pair.user_data_["trades"] = pd.DataFrame(columns=self.TRADES_COLUMNS)
+        pair.user_data_["trades"] = pd.DataFrame(columns=self.TRADES_COLUMNS)  # type: ignore
         pair.user_data_["is_cointegrated"] = False
 
         open_threshold = config["dis-equilibrium_open_trshld"]
@@ -357,7 +363,7 @@ class SlidingFitStrategy(PairsTradingStrategy):
         ]
         return pd.DataFrame(
             trd_signal_tuples,
-            columns=self.TRADES_COLUMNS,
+            columns=self.TRADES_COLUMNS,  # type: ignore
         )
 
     def _get_close_trades(self, pair: TradingPair, close_threshold: float) -> Optional[pd.DataFrame]:
@@ -404,9 +410,11 @@ class SlidingFitStrategy(PairsTradingStrategy):
         # Add tuples to data frame
         return pd.DataFrame(
             trd_signal_tuples,
-            columns=self.TRADES_COLUMNS,
+            columns=self.TRADES_COLUMNS,  # type: ignore
         )
 
+    def reset(self):
+        self.curr_training_start_idx_ = 0