Update utils/plotting.py

utils/plotting.py

@@ -3,12 +3,13 @@ import matplotlib.dates as mdates
 
 def plot_stock_data_with_signals(stock_data):
     """
-    Plots the stock data along with SMAs, Bollinger Bands, and buy/sell signals.
+    Creates a plot of stock data along with SMAs, Bollinger Bands, and buy/sell signals,
+    tailored for display in a Streamlit app.
 
     Parameters:
-    - stock_data (pd.DataFrame): The stock data with 'Close', 'SMA_21', 'SMA_50', 'BB_Upper', 'BB_Lower', 'Buy_Signal', and 'Sell_Signal' columns.
+    - stock_data (pd.DataFrame): The stock data with 'Close', 'SMA_21', 'SMA_50',
+      'BB_Upper', 'BB_Lower', 'Buy_Signal', and 'Sell_Signal' columns.
     """
-    # Setting up the plot
     fig, ax = plt.subplots(figsize=(14, 7))
 
     # Plotting the closing prices
@@ -22,28 +23,24 @@ def plot_stock_data_with_signals(stock_data):
     ax.plot(stock_data.index, stock_data['BB_Upper'], label='Upper Bollinger Band', color='red', linestyle='--', alpha=0.5)
     ax.plot(stock_data.index, stock_data['BB_Lower'], label='Lower Bollinger Band', color='cyan', linestyle='--', alpha=0.5)
 
-    # Highlighting buy signals
+    # Highlighting buy and sell signals
     buy_signals = stock_data[stock_data['Buy_Signal']]
-    ax.scatter(buy_signals.index, buy_signals['Close'], label='Buy Signal', marker='^', color='green', alpha=1, s=100)
-
-    # Highlighting sell signals
     sell_signals = stock_data[stock_data['Sell_Signal']]
+    ax.scatter(buy_signals.index, buy_signals['Close'], label='Buy Signal', marker='^', color='green', alpha=1, s=100)
     ax.scatter(sell_signals.index, sell_signals['Close'], label='Sell Signal', marker='v', color='red', alpha=1, s=100)
 
-    # Beautifying the plot
+    # Setting title and labels
     ax.set_title("Stock Price with Indicators and Signals")
     ax.set_xlabel("Date")
     ax.set_ylabel("Price")
-    ax.legend()
 
-    # Formatting the x-axis to show dates nicely
+    # Formatting date on the x-axis
     ax.xaxis.set_major_locator(mdates.WeekdayLocator())
     ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
-    fig.autofmt_xdate() # Rotate date labels to prevent overlap
+    fig.autofmt_xdate()
 
-    plt.show()
+    # Adding legend
+    ax.legend()
 
-if __name__ == "__main__":
-    # Example usage:
-    # Generate or fetch your stock_data DataFrame with necessary columns before calling this function
-    pass
+    # Instead of plt.show(), just return the figure object
+    return fig