Source code for richy.transactions.dividends

import logging
from datetime import date

from ..core.models import User
from .models import (
    EtfDividendTransaction,
    ShareDividendTransaction,
    Transaction,
)
from .transactions import Transactions

LOGGER = logging.getLogger(__name__)




class DividendTransactions:
    """
    This class handles ``ShareDividendTransaction`` or ``EtfDividendTransaction``
    model creation/update. It calculates based on ``Transaction`` (currently owned
    stocks) and ``shares.Dividend`` (dividend data) transactions that are income
    for share holder.
    """

    def __init__(self, *, share=None, etf=None):
        """
        :param Share share. Share model instance.
        :param Etf etf. Etf model instance.
        """

        if share is None and etf is None:
            raise Exception("share or etf parameters must be provided.")

        if share:
            self.item = share
            self.df_item_type = "share"
            self.dividend_model = ShareDividendTransaction
        else:
            self.item = etf
            self.df_item_type = "etf"
            self.dividend_model = EtfDividendTransaction



    def calculate(self):
        """
        Calculates (creates/updates) ``ShareDividendTransaction`` or ``EtfDividendTransaction``
        model records based on owned stocks now and in past.

        Also recalculates/removes already created transaction in the past.
        """

        # Fetch only those users that have at least one transaction.
        for user in User.objects.filter(
            pk__in=Transaction.objects.distinct("user").values_list("user", flat=True)
        ):
            dividends = self.item.dividend_set.filter(
                record_date__lte=date.today()
            ).order_by("record_date")
            df = Transactions(user).get_transaction_basic_stats()

            if df.empty:
                LOGGER.debug(
                    "Empty dataframe (no transactions) - terminating dividend calculation."
                )

                return

            df = df[df["type"] == self.df_item_type]

            for div in dividends:
                df_filtered = df[df["date"] <= div.payment_date]
                df_filtered = df_filtered[df_filtered["symbol"] == self.item.symbol]

                # If no shares were held at the time dividend ex date we skip it.
                if df_filtered.empty:
                    continue

                # Calculate currently (div.payment_date) owned shares.
                df_filtered_uniq = df_filtered.groupby(df_filtered.index).first()
                owned_shares = df_filtered_uniq.amount.sum()

                # If some were owned.
                if owned_shares > 0:
                    record, created = self.dividend_model.objects.update_or_create(
                        user=user,
                        dividend=div,
                        defaults={
                            "shares": owned_shares,
                            "amount": div.amount_value * owned_shares,
                        },
                    )
                    record.transactions.set(df_filtered_uniq.index.to_list())

                    if created:
                        LOGGER.debug(
                            "Dividend transaction has been calculated for "
                            f"{self.item} for date {div.payment_date} for user with ID {user.pk}."
                        )
                    else:
                        LOGGER.debug(
                            "Dividend transaction has been recalculated for "
                            f"{self.item} for date {div.payment_date} for user with ID {user.pk}."
                        )
                else:
                    try:
                        self.dividend_model.objects.by_user(user).get(
                            dividend=div
                        ).delete()

                        LOGGER.debug(
                            f"Dividend transaction for share {self.item.symbol} "
                            f"and {div.payment_date} for user with ID {user.pk} has been deleted."
                        )
                    except self.dividend_model.DoesNotExist:
                        pass