REF: lib.resample_apply() to better handle Series/DataFrame

Author: kernc

backtesting/lib.py

@@ -133,7 +133,7 @@ def quantile(series, quantile=None):
 
 
 def resample_apply(rule: str,
-                   func: Callable,
+                   func: Optional[Callable[..., Sequence]],
                    series,
                    *args,
                    agg='last',
@@ -205,6 +205,10 @@ def SMA(series, n):
                 self.sma = self.I(SMA, daily, 10, plot=False)
 
     """
+    if func is None:
+        def func(x, *_, **__):
+            return x
+
     if not isinstance(series, (pd.Series, pd.DataFrame)):
         assert isinstance(series, _Array), \
             'resample_apply() takes either a `pd.Series`, `pd.DataFrame`, ' \
@@ -228,14 +232,20 @@ def SMA(series, n):
         def strategy_I(func, *args, **kwargs):
             return func(*args, **kwargs)
 
-    # Resample back to data index
     def wrap_func(resampled, *args, **kwargs):
-        ind = pd.Series(np.asarray(func(resampled, *args, **kwargs)),
-                        index=resampled.index,
-                        name=resampled.name)
-        ind = ind.reindex(index=series.index | resampled.index,
-                          method='ffill').reindex(series.index)
-        return ind
+        result = func(resampled, *args, **kwargs)
+        if not isinstance(result, pd.DataFrame) and not isinstance(result, pd.Series):
+            result = np.asarray(result)
+            if result.ndim == 1:
+                result = pd.Series(result, name=resampled.name)
+            elif result.ndim == 2:
+                result = pd.DataFrame(result.T)
+        # Resample back to data index
+        if not result.index.is_all_dates:
+            result.index = resampled.index
+        result = result.reindex(index=series.index | resampled.index,
+                                method='ffill').reindex(series.index)
+        return result
 
     wrap_func.__name__ = func.__name__
 

backtesting/test/_test.py

@@ -530,6 +530,12 @@ def resets_index(*args):
         self.assertTrue((res.dropna() == res2.dropna()).all())
         self.assertTrue((res.index == res2.index).all())
 
+        res3 = resample_apply('D', None, EURUSD)
+        self.assertIn('Volume', res3)
+
+        res3 = resample_apply('D', lambda df: (df.Close, df.Close), EURUSD)
+        self.assertIsInstance(res3, pd.DataFrame)
+
     def test_plot_heatmaps(self):
         bt = Backtest(GOOG, SmaCross)
         stats, heatmap = bt.optimize(fast=range(2, 7, 2),