DividendFlow

dividend.py

import bt
import pandas as pd


class DividendFlow(bt.Algo):
    
    def __init__(self, dividends: pd.DataFrame):
        self.dividends = dividends
        self.initialized = False
        self.last_date = None
        self.payments = {}
        self.windows = set()
        super(DividendFlow, self).__init__()
        
    def __call__(self, target: bt.Strategy):
        # skip first execution where we don't have anything allocated, so no dividends to reinvest
        if not self.initialized:
            self.initialized = True
            self.last_date = target.now
            return True
        
        start_window = self.last_date.strftime('%Y-%m-%d')
        end_window = target.now.strftime('%Y-%m-%d')
        dividends_range: pd.DataFrame = self.dividends.loc[start_window:end_window]
        self.last_date = target.now
        
        for date, loc in dividends_range.iterrows():
            if date in self.windows:
                continue
            else:
                self.windows.add(date)

            for ticker, dividend in loc.items():
                try:
                    stock_amount = target.positions[ticker][date]  # check if we had this stock at this date
                except KeyError:
                    continue
                
                try:
                    cash = dividend * stock_amount
                except KeyError:
                    break
                
                if cash > 0:
                    target.adjust(cash, flow=False)

        return True