Change indicators to return their results and work on three indicator strat

strategy_evaluation/indicators.py

@@ -17,9 +17,13 @@ def bollinger_band(df, symbol, period=20, m=2):
     std = df[symbol].rolling(period).std()
     boll_up = boll_sma + m * std
     boll_lo = boll_sma - m * std
-    df[f"{symbol}-Boll({period})-sma"] = boll_sma
-    df[f"{symbol}-Boll({period})-up"] = boll_up
-    df[f"{symbol}-Boll({period})-lo"] = boll_lo
+    key_sma = f"{symbol}-Boll({period})-sma"
+    key_up = f"{symbol}-Boll({period})-up"
+    key_lo = f"{symbol}-Boll({period})-lo"
+    df[key_sma] = boll_sma
+    df[key_up] = boll_up
+    df[key_lo] = boll_lo
+    return df[[key_sma, key_up, key_lo]]
 
 
 def sma(df, symbol, period):
@@ -35,14 +39,28 @@ def sma(df, symbol, period):
 
 
 def ema(df, symbol, period):
-    """Adds a new column to the dataframe EMA(period)"""
-    df[f"{symbol}-ema({period})"] = df[symbol].ewm(span=period).mean()
+    """Adds EMA for one or multiple periods to df and returns EMAs"""
+    if type(period) is int:
+        period = [period]
+    keys = []
+    for p in period:
+        key = f"{symbol}-ema({p})"
+        df[key] = df[symbol].ewm(span=period).mean()
+        keys.append(key)
+    return df[keys]
 
 
 def price_sma(df, symbol, period):
     """Calculates SMA and adds new column price divided by SMA to the df."""
-    sma = df[symbol].rolling(period).mean()
-    df[f"{symbol}-price/sma({period})"] = df[symbol] / sma
+    if type(period) is int:
+        period = [period]
+    keys = []
+    for p in period:
+        key = f"{symbol}-price/sma({p})"
+        sma = df[symbol].rolling(p).mean()
+        df[key] = df[symbol] / sma
+        keys.append(key)
+    return df[keys]
 
 
 def rsi(df, symbol, period=14):
@@ -61,7 +79,7 @@ def rsi(df, symbol, period=14):
     # Add one to get 'period' price changes (first change is nan).
     period += 1
     df[key] = df[symbol].rolling(period).apply(rsi)
-    return df[key]
+    return df[[key]]
 
 
 def macd(df, symbol):