ENH: Add pandas accessors .df and .s

Author: kernc

backtesting/_util.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Sequence
 from numbers import Number
 
@@ -68,6 +69,11 @@ def __float__(self):
             return super().__float__()
 
     def to_series(self):
+        warnings.warn("`.to_series()` is deprecated. For pd.Series conversion, use accessor `.s`")
+        return self.s
+
+    @property
+    def s(self) -> pd.Series:
         return pd.Series(self, index=self._opts['data'].index, name=self.name)
 
 
@@ -82,15 +88,13 @@ class _Data:
     and the returned "series" are _not_ `pd.Series` but `np.ndarray`
     for performance reasons.
     """
-    def __init__(self, df):
+    def __init__(self, df: pd.DataFrame):
+        self.__df = df
         self.__i = len(df)
         self.__pip = None
         self.__cache = {}
-
-        self.__arrays = {col: _Array(arr, data=self)
-                         for col, arr in df.items()}
-        # Leave index as Series because pd.Timestamp nicer API to work with
-        self.__arrays['__index'] = df.index.copy()
+        self.__arrays = None
+        self._update_arrays()
 
     def __getitem__(self, item):
         return getattr(self, item)
@@ -105,17 +109,29 @@ def _set_length(self, i):
         self.__i = i
         self.__cache.clear()
 
+    def _update_arrays(self):
+        self.__arrays = {col: _Array(arr, data=self)
+                         for col, arr in self.__df.items()}
+        # Leave index as Series because pd.Timestamp nicer API to work with
+        self.__arrays['__index'] = self.__df.index.copy()
+
     def __len__(self):
         return self.__i
 
+    @property
+    def df(self) -> pd.DataFrame:
+        return (self.__df.iloc[:self.__i]
+                if self.__i < len(self.__df)
+                else self.__df)
+
     @property
     def pip(self):
         if self.__pip is None:
             self.__pip = 10**-np.median([len(s.partition('.')[-1])
                                          for s in self.__arrays['Close'].astype(str)])
         return self.__pip
 
-    def __get_array(self, key):
+    def __get_array(self, key) -> _Array:
         arr = self.__cache.get(key)
         if arr is None:
             arr = self.__cache[key] = self.__arrays[key][:self.__i]
@@ -142,8 +158,8 @@ def Volume(self):
         return self.__get_array('Volume')
 
     @property
-    def index(self):
-        return self.__get_array('__index')
+    def index(self) -> pd.Index:
+        return self.__get_array('__index')  # type: ignore
 
     # Make pickling in Backtest.optimize() work with our catch-all __getattr__
     def __getstate__(self):

backtesting/backtesting.py

@@ -969,11 +969,12 @@ def run(self, **kwargs) -> pd.Series:
 
         Keyword arguments are interpreted as strategy parameters.
         """
-        data = _Data(self._data)
+        data = _Data(self._data.copy(deep=False))
         broker = self._broker(data=data)  # type: _Broker
         strategy = self._strategy(broker, data, kwargs)  # type: Strategy
 
         strategy.init()
+        data._update_arrays()  # Strategy.init might have changed/added to data.df
 
         # Indicators used in Strategy.next()
         indicator_attrs = {attr: indicator

backtesting/lib.py

@@ -209,7 +209,7 @@ def SMA(series, n):
         assert isinstance(series, _Array), \
             'resample_apply() takes either a `pd.Series`, `pd.DataFrame`, ' \
             'or a `Strategy.data.*` array'
-        series = series.to_series()
+        series = series.s
 
     series = series.copy()  # XXX: pandas 1.0.1 bug https://github.com/pandas-dev/pandas/issues/31710  # noqa: E501
     resampled = series.resample(rule, label='right').agg(agg).dropna()

backtesting/test/_test.py

@@ -555,7 +555,7 @@ def test_plot_heatmaps(self):
     def test_SignalStrategy(self):
         class S(SignalStrategy):
             def init(self):
-                sma = self.data.Close.to_series().rolling(10).mean()
+                sma = self.data.Close.s.rolling(10).mean()
                 self.set_signal(self.data.Close > sma,
                                 self.data.Close < sma)
 
@@ -568,7 +568,7 @@ def init(self):
                 super().init()
                 self.set_atr_periods(40)
                 self.set_trailing_sl(3)
-                self.sma = self.I(lambda: self.data.Close.to_series().rolling(10).mean())
+                self.sma = self.I(lambda: self.data.Close.s.rolling(10).mean())
 
             def next(self):
                 super().next()
@@ -596,6 +596,22 @@ class Class:
         for s in ('Open', 'High', 'Low', 'Close'):
             self.assertEqual(_as_str(_Array([1], name=s)), s[0])
 
+    def test_pandas_accessors(self):
+        class S(Strategy):
+            def init(self):
+                close, index = self.data.Close, self.data.index
+                assert close.s.equals(pd.Series(close, index=index))
+                assert self.data.df['Close'].equals(pd.Series(close, index=index))
+                self.data.df['new_key'] = 2 * close
+
+            def next(self):
+                close, index = self.data.Close, self.data.index
+                assert close.s.equals(pd.Series(close, index=index))
+                assert self.data.df['Close'].equals(pd.Series(close, index=index))
+                assert self.data.df['new_key'].equals(pd.Series(self.data.new_key, index=index))
+
+        Backtest(GOOG.iloc[:20], S).run()
+
 
 @unittest.skipUnless(
     os.path.isdir(os.path.join(os.path.dirname(__file__),