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"{self.symbol}" else: return f"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, )