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,
)