market/trading/models.py

from datetime import date
from typing import Mapping, Optional, Sequence
from django.conf import settings
from django.contrib.auth.models import AbstractBaseUser, PermissionsMixin
from django.db import models
from django.db.models.signals import post_save, pre_save
from django.dispatch import receiver
from django.utils.translation import gettext_lazy as _
from frozendict import frozendict
from hashid_field import HashidAutoField
from trading.managers import UserManager


class User(PermissionsMixin, AbstractBaseUser):
    """
    Base user model, with auth information already filled in.
    """

    username = models.CharField(max_length=30, unique=True)
    email = models.EmailField(unique=True)
    is_superuser = models.BooleanField(default=False)
    is_active = models.BooleanField(default=True)
    date_joined = models.DateTimeField(_("date joined"), auto_now_add=True)

    max_commodities_allowed = models.PositiveIntegerField(
        default=settings.TRADING_FREE_COMMODITIES,
        null=True,
        verbose_name=_("number of commodities this user is allowed to create"),
    )

    objects = UserManager()

    USERNAME_FIELD = "username"
    EMAIL_FIELD = "email"
    REQUIRED_FIELDS = (EMAIL_FIELD,)

    class Meta:
        verbose_name = _("user")
        verbose_name_plural = _("users")

    def get_full_name(self):
        """
        Returns the first_name plus the last_name, with a space in between.
        """
        return self.username

    def get_short_name(self):
        """
        Returns the short name for the user.
        """
        return self.username

    def email_user(self, subject, message, from_email=None, **kwargs):
        """
        Sends an email to this User.
        """
        send_mail(subject, message, from_email, [self.email], **kwargs)

    def can_create_commodity(self) -> bool:
        if self.max_commodities_allowed is None:
            return True
        return self.max_commodities_allowed > self.commodity_set.count()

    def unused_commodities(self) -> Optional[int]:
        if self.max_commodities_allowed is None:
            return None
        return self.max_commodities_allowed - self.commodity_set.count()

    def unused_invites(self) -> Sequence["Invite"]:
        return self.invite_set.filter(accepted_user__isnull=True)

    def balance_of(self, commodity: "Commodity") -> int:
        return self.balances().get(commodity, 0)

    def balances(self) -> Mapping["Commodity", int]:
        from django.db.models import Sum

        credits = {
            row["commodity"]: row["amount__sum"]
            for row in self.credits.values("commodity").annotate(Sum("amount"))
        }
        debits = {
            row["commodity"]: row["amount__sum"]
            for row in self.debits.values("commodity").annotate(Sum("amount"))
        }
        ipos = {
            row["pk"]: row["in_circulation"]
            for row in self.commodity_set.values("pk", "in_circulation")
        }
        keys = set(ipos.keys()) | set(credits.keys()) | set(debits.keys())
        return frozendict(
            [
                (
                    Commodity.objects.get(pk=pk),
                    credits.get(pk, 0) + ipos.get(pk, 0) - debits.get(pk, 0),
                )
                for pk in keys
            ]
        )


@receiver(post_save, sender=User)
def _user_post_save(sender, instance, created, **kwargs):
    # Create invites for the user, if they are permitted invites
    if created and settings.TRADING_FREE_INVITES:
        for _ in range(settings.TRADING_FREE_INVITES):
            Invite.objects.create(inviter=instance).save()


class Invite(models.Model):
    """
    An invitation that is given out to a user.
    """

    id = HashidAutoField(primary_key=True)
    inviter = models.ForeignKey(User, on_delete=models.SET_NULL, null=True)
    accepted_user = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        null=True,
        default=None,
        related_name="+",
        verbose_name=_("accepted user"),
    )

    class Meta:
        verbose_name = _("invite")
        verbose_name_plural = _("invites")


class Commodity(models.Model):
    """
    A commodity with some kind of market cap.

    A commodity has an associated "owner" account, which determines where the IPO for commodity
    goes.
    """

    id = HashidAutoField(primary_key=True)
    name = models.CharField(
        blank=False, unique=True, max_length=100, verbose_name=_("commodity name")
    )
    in_circulation = models.PositiveIntegerField()
    created_by = models.ForeignKey(User, null=True, on_delete=models.SET_NULL)
    symbol = models.CharField(
        blank=True, unique=True, max_length=6, verbose_name=_("symbol")
    )

    class Meta:
        verbose_name = _("commodity")
        verbose_name_plural = _("commodities")

    def get_absolute_url(self):
        return "/c/detail/%s" % self.pk

    def html_symbol(self) -> str:
        if self.symbol:
            return f"<span>{self.symbol}</span>"
        else:
            return f"<span>units of {self.name}"


class MaxCommodityError(Exception):
    def __init__(self, user: User):
        self.user = user
        super().__init__("maximum commodity count reached")


class Tx(models.Model):
    """
    A transaction from one account to another.
    """

    id = HashidAutoField(primary_key=True)
    instant = models.DateTimeField(auto_now_add=True, editable=False)
    source = models.ForeignKey(
        User,
        on_delete=models.PROTECT,
        related_name="debits",
        related_query_name="debit_tx",
        verbose_name=_("source"),
        editable=False,
    )
    dest = models.ForeignKey(
        User,
        on_delete=models.PROTECT,
        related_name="credits",
        related_query_name="credit_tx",
        verbose_name=_("destination"),
        editable=False,
    )
    amount = models.PositiveIntegerField(verbose_name=_("amount"), editable=False)
    commodity = models.ForeignKey(Commodity, on_delete=models.PROTECT, editable=False)

    class Meta:
        verbose_name = _("transaction")
        verbose_name_plural = _("transactions")


class TxRequests(models.Model):
    """
    A transaction request between two users.
    """

    WAITING = "WAIT"
    DECLINED = "DECL"
    REVISED = "REV"

    STATUS_CHOICES = [
        (WAITING, "Waiting"),
        (DECLINED, "Declined"),
        (REVISED, "Revised"),
    ]

    status = models.CharField(
        max_length=max([len(c) for c in STATUS_CHOICES]),
        choices=STATUS_CHOICES,
        default=WAITING,
    )
    previous_request = models.OneToOneField(
        "TxRequests",
        null=True,
        default=None,
        on_delete=models.SET_NULL,
        related_name="revised_request",
    )

    source = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        null=True,
        related_name="transaction_requests_sent",
    )
    dest = models.ForeignKey(
        User,
        on_delete=models.CASCADE,
        null=True,
        related_name="transaction_requests_received",
    )

    source_sends = models.ForeignKey(
        Commodity, on_delete=models.CASCADE, related_name="+", null=True,
    )
    dest_sends = models.ForeignKey(
        Commodity, on_delete=models.CASCADE, related_name="+", null=True,
    )

    source_amount = models.PositiveIntegerField()
    dest_amount = models.PositiveIntegerField()


class BalanceError(Exception):
    def __init__(
        self, balance: int, source: User, dest: User, amount: int, commodity: Commodity
    ) -> None:
        self.balance = balance
        self.source = source
        self.dest = dest
        self.amount = amount
        self.commodity = commodity
        super().__init__("balance too low to create transaction")


@receiver(pre_save, sender=Tx)
def _tx_pre_save(sender, instance, *args, **kwargs):
    """
    Ensures that a transaction cannot be created if the source user doesn't have enough of the given
    commodity in their account.
    """
    if instance.pk is None:
        balance = instance.source.balance_of(instance.commodity)
        if instance.amount > balance:
            raise BalanceError(
                balance,
                instance.source,
                instance.dest,
                instance.amount,
                instance.commodity,
            )