vote.py 3.66 KB
Newer Older
luxcem's avatar
luxcem committed
1 2 3 4 5
from django.db import models
from django.db.models import Sum, F
from django.contrib.contenttypes.fields import GenericForeignKey
from django.contrib.contenttypes.fields import GenericRelation
from django.contrib.contenttypes.models import ContentType
6
from django.utils.translation import ugettext_lazy as _
luxcem's avatar
luxcem committed
7 8 9 10 11 12 13


class UnDVotes(models.Model):
    """
    Up and down vote model
    """

14 15 16 17 18
    #: username
    username = models.CharField(max_length=255, null=False)

    #: Score of the vote
    score = models.IntegerField(default=1)
luxcem's avatar
luxcem committed
19 20 21 22 23 24
    # Django generic relation
    content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE)
    object_id = models.PositiveIntegerField()
    #: The Voted object
    content_object = GenericForeignKey("content_type", "object_id")

25 26 27 28
    class Meta:
        verbose_name = _("Vote")
        verbose_name_plural = _("Votes")

luxcem's avatar
luxcem committed
29
    def __str__(self):
luxcem's avatar
luxcem committed
30 31
        return "{}:{}:{}".format(
            self.username, self.content_object, self.score)
luxcem's avatar
luxcem committed
32 33 34 35 36 37 38 39 40 41


class UnDVotedMixin(models.Model):
    """
    A mixin to attach to a model that has up and down votes
    """
    #: Votes
    und_votes = GenericRelation(UnDVotes)
    #: Score of the model
    und_score = models.IntegerField(default=0)
cynddl's avatar
cynddl committed
42 43
    und_score_up = models.IntegerField(default=0)
    und_score_down = models.IntegerField(default=0)
luxcem's avatar
luxcem committed
44

45 46 47
    class Meta:
        abstract = True

48
    def upvote(self, username):
luxcem's avatar
luxcem committed
49 50 51
        diff_score = 0
        try:
            # Already voted content
52
            vote = self.und_votes.get(username=username)
luxcem's avatar
luxcem committed
53 54 55 56 57 58 59 60 61 62
            if vote.score == 1:
                # Cancel previous upvote
                vote.delete()
                diff_score = -1
            else:
                # Previously downvoted
                vote.score = 1
                vote.save()
                diff_score = 2
        except:
63
            vote = UnDVotes(content_object=self, username=username, score=1)
luxcem's avatar
luxcem committed
64 65 66
            vote.save()
            diff_score = 1

luxcem's avatar
luxcem committed
67
        self.und_score += diff_score
cynddl's avatar
cynddl committed
68
        self.und_score_up += diff_score
luxcem's avatar
luxcem committed
69
        # Update self score, use update and filter to avoid triggering signals
luxcem's avatar
luxcem committed
70
        self.__class__.objects.filter(id=self.id).update(
luxcem's avatar
luxcem committed
71 72 73
            und_score=F("und_score") + diff_score,
            und_score_up=F("und_score_up") + diff_score
        )
luxcem's avatar
luxcem committed
74

75
    def downvote(self, username):
luxcem's avatar
luxcem committed
76 77 78
        diff_score = 0
        try:
            # Already voted content
79
            vote = self.und_votes.get(username=username)
luxcem's avatar
luxcem committed
80 81 82 83 84 85 86 87 88 89
            if vote.score == -1:
                # Cancel previous downvote
                vote.delete()
                diff_score = 1
            else:
                # Previously upvoted
                vote.score = -1
                vote.save()
                diff_score = -2
        except:
90
            vote = UnDVotes(content_object=self, username=username, score=-1)
luxcem's avatar
luxcem committed
91 92 93
            vote.save()
            diff_score = -1

luxcem's avatar
luxcem committed
94
        self.und_score += diff_score
cynddl's avatar
cynddl committed
95
        self.und_score_down += diff_score
luxcem's avatar
luxcem committed
96 97
        # Update self score, use update and filter to avoid triggering signals
        self.__class__.objects.filter(id=self.id).update(
luxcem's avatar
luxcem committed
98 99 100
            und_score=F("und_score") + diff_score,
            und_score_down=F("und_score_down") + diff_score
        )
luxcem's avatar
luxcem committed
101 102 103

    def update_und_score(self):
        """Reset score to the correct count (should not be necessary)"""
cynddl's avatar
cynddl committed
104 105 106 107 108 109
        score_up = self.und_votes.filter(score=1).count()
        score_down = self.und_votes.filter(score=-1).count()
        score = score_up + score_down

        self.und_score_up = score_up
        self.und_score_down = score_down
luxcem's avatar
luxcem committed
110
        self.und_score = score
luxcem's avatar
luxcem committed
111 112
        # Update self score, use update and filter to avoid triggering signals
        self.__class__.objects.filter(id=self.id).update(und_score=score)