latch_adjust_dataframe.py

def latch_adjusted_dataframe(input_dataframe, shift_period):
    df = input_dataframe.copy()

    # latch-adjust returns
    df.reset_index(inplace=True)

    df['buy_signal'] = df['buy_signal'] & (df['buy_signal'].shift(1) == False)
    df['sell_signal'] = df['sell_signal'] & (df['sell_signal'].shift(1) == False)

    # not specifying a condition, because we are going to modify the same condition.
    buy_lockout = 0
    sell_lockout = 0
    for row in df.itertuples():
        if row.buy_signal and buy_lockout == 0:
            buy_lockout = shift_period
        elif row.buy_signal and buy_lockout > 0:
            df.loc[row[0], 'buy_signal'] = False
            buy_lockout -= 1
        else:
            buy_lockout -= 1

        if row.sell_signal and sell_lockout == 0:
            sell_lockout = shift_period
        elif row.sell_signal and sell_lockout > 0:
            df.loc[row[0], 'sell_signal'] = False
            sell_lockout -= 1
        else:
            sell_lockout -= 1

        if buy_lockout < 0:
            buy_lockout = 0
        if sell_lockout < 0:
            sell_lockout = 0

    df['buy_block'] = False
    df['sell_block'] = False

    for row in df[df['buy_signal'] == True].itertuples():
        index = row[0]
        df.loc[index+1:index+shift_period, 'buy_block'] = True
        start_index = index-shift_period if index-shift_period >= 0 else 0
        df.loc[start_index:index-1, 'buy_block'] = True

    for row in df[df['sell_signal'] == True].itertuples():
        index = row[0]
        df.loc[index+1:index+shift_period, 'sell_block'] = True
        start_index = index-shift_period if index-shift_period >= 0 else 0
        df.loc[start_index:index-1, 'sell_block'] = True

    df.set_index('index', inplace=True)
    
    return df