indicators/rsi.py

import pandas as pd

def calculate_rsi(prices, length=14, oversold=30, overbought=70):
    """
    Calculates the Relative Strength Index (RSI) and marks the oversold and overbought conditions.
    
    Parameters:
    - prices (pd.Series): A pandas Series containing the stock's closing prices.
    - length (int): The length of the RSI period. Defaults to 14.
    - oversold (int): The level at which the asset is considered oversold. Defaults to 30.
    - overbought (int): The level at which the asset is considered overbought. Defaults to 70.
    
    Returns:
    - pd.DataFrame: A DataFrame containing the RSI values, and flags for oversold and overbought conditions.
    """
    delta = prices.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=length).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=length).mean()

    rs = gain / loss
    rsi = 100 - (100 / (1 + rs))

    rsi_df = pd.DataFrame(data={'RSI': rsi})
    rsi_df['Oversold'] = rsi_df['RSI'] < oversold
    rsi_df['Overbought'] = rsi_df['RSI'] > overbought

    return rsi_df

if __name__ == "__main__":
    # Example usage
    data = {'Close': [22, 24, 23, 25, 26, 28, 27, 29, 30, 32, 31, 33]}
    prices = pd.Series(data['Close'])
    
    # User-defined parameters for RSI
    length = 14  # RSI period
    oversold = 30  # Oversold threshold
    overbought = 70  # Overbought threshold
    
    # Calculate RSI
    rsi_df = calculate_rsi(prices, length, oversold, overbought)
    
    print(rsi_df)