Fix pagevote and bookmark bug for new objects

judge/models/bookmark.py

@@ -53,3 +53,13 @@ class MakeBookMark(models.Model):
         unique_together = ["user", "bookmark"]
         verbose_name = _("make bookmark")
         verbose_name_plural = _("make bookmarks")
+
+
+class Bookmarkable:
+    def get_or_create_bookmark(self):
+        if self.bookmark.count():
+            return self.bookmark.first()
+        new_bookmark = BookMark()
+        new_bookmark.linked_object = self
+        new_bookmark.save()
+        return new_bookmark

judge/models/contest.py

@@ -22,6 +22,8 @@ from judge.models.problem import Problem
 from judge.models.profile import Organization, Profile
 from judge.models.submission import Submission
 from judge.ratings import rate_contest
+from judge.models.pagevote import PageVotable
+from judge.models.bookmark import Bookmarkable
 
 __all__ = [
     "Contest",
@@ -75,7 +77,7 @@ class ContestTag(models.Model):
         verbose_name_plural = _("contest tags")
 
 
-class Contest(models.Model):
+class Contest(models.Model, PageVotable, Bookmarkable):
     SCOREBOARD_VISIBLE = "V"
     SCOREBOARD_AFTER_CONTEST = "C"
     SCOREBOARD_AFTER_PARTICIPATION = "P"

judge/models/interface.py

@@ -11,8 +11,8 @@ from mptt.fields import TreeForeignKey
 from mptt.models import MPTTModel
 
 from judge.models.profile import Organization, Profile
-from judge.models.pagevote import PageVote
-from judge.models.bookmark import BookMark
+from judge.models.pagevote import PageVotable
+from judge.models.bookmark import Bookmarkable
 
 __all__ = ["MiscConfig", "validate_regex", "NavigationBar", "BlogPost"]
 
@@ -74,7 +74,7 @@ class NavigationBar(MPTTModel):
             return pattern
 
 
-class BlogPost(models.Model):
+class BlogPost(models.Model, PageVotable, Bookmarkable):
     title = models.CharField(verbose_name=_("post title"), max_length=100)
     authors = models.ManyToManyField(Profile, verbose_name=_("authors"), blank=True)
     slug = models.SlugField(verbose_name=_("slug"))
@@ -132,22 +132,6 @@ class BlogPost(models.Model):
             and self.authors.filter(id=user.profile.id).exists()
         )
 
-    def get_or_create_pagevote(self):
-        if self.pagevote.count():
-            return self.pagevote.first()
-        new_pagevote = PageVote()
-        new_pagevote.linked_object = self
-        new_pagevote.save()
-        return new_pagevote
-
-    def get_or_create_bookmark(self):
-        if self.bookmark.count():
-            return self.bookmark.first()
-        new_bookmark = BookMark()
-        new_bookmark.linked_object = self
-        new_bookmark.save()
-        return new_bookmark
-
     class Meta:
         permissions = (("edit_all_post", _("Edit all posts")),)
         verbose_name = _("blog post")

judge/models/pagevote.py

@@ -48,3 +48,13 @@ class PageVoteVoter(models.Model):
         unique_together = ["voter", "pagevote"]
         verbose_name = _("pagevote vote")
         verbose_name_plural = _("pagevote votes")
+
+
+class PageVotable:
+    def get_or_create_pagevote(self):
+        if self.pagevote.count():
+            return self.pagevote.first()
+        new_pagevote = PageVote()
+        new_pagevote.linked_object = self
+        new_pagevote.save()
+        return new_pagevote

judge/models/problem.py

@@ -14,8 +14,8 @@ from django.utils.functional import cached_property
 from django.utils.translation import gettext_lazy as _
 
 from judge.fulltext import SearchQuerySet
-from judge.models.pagevote import PageVote
-from judge.models.bookmark import BookMark
+from judge.models.pagevote import PageVotable
+from judge.models.bookmark import Bookmarkable
 from judge.models.profile import Organization, Profile
 from judge.models.runtime import Language
 from judge.user_translations import gettext as user_gettext
@@ -124,7 +124,7 @@ class TranslatedProblemQuerySet(SearchQuerySet):
         )
 
 
-class Problem(models.Model):
+class Problem(models.Model, PageVotable, Bookmarkable):
     code = models.CharField(
         max_length=20,
         verbose_name=_("problem code"),
@@ -549,22 +549,6 @@ class Problem(models.Model):
         cache.set(key, result)
         return result
 
-    def get_or_create_pagevote(self):
-        if self.pagevote.count():
-            return self.pagevote.first()
-        new_pagevote = PageVote()
-        new_pagevote.linked_object = self
-        new_pagevote.save()
-        return new_pagevote
-
-    def get_or_create_bookmark(self):
-        if self.bookmark.count():
-            return self.bookmark.first()
-        new_bookmark = BookMark()
-        new_bookmark.linked_object = self
-        new_bookmark.save()
-        return new_bookmark
-
     def save(self, *args, **kwargs):
         super(Problem, self).save(*args, **kwargs)
         if self.code != self.__original_code:
@@ -669,7 +653,7 @@ class LanguageTemplate(models.Model):
         verbose_name_plural = _("language-specific templates")
 
 
-class Solution(models.Model):
+class Solution(models.Model, PageVotable, Bookmarkable):
     problem = models.OneToOneField(
         Problem,
         on_delete=CASCADE,

judge/views/bookmark.py

@@ -71,6 +71,5 @@ def undobookmark_page(request):
 class BookMarkDetailView(TemplateResponseMixin, SingleObjectMixin, View):
     def get_context_data(self, **kwargs):
         context = super(BookMarkDetailView, self).get_context_data(**kwargs)
-        queryset = self.object.bookmark
-        context["bookmark"] = queryset.first()
+        context["bookmark"] = self.object.get_or_create_bookmark()
         return context

judge/views/pagevote.py

@@ -101,6 +101,5 @@ class PageVoteDetailView(TemplateResponseMixin, SingleObjectMixin, View):
 
     def get_context_data(self, **kwargs):
         context = super(PageVoteDetailView, self).get_context_data(**kwargs)
-        queryset = self.object.pagevote
-        context["pagevote"] = queryset.first()
+        context["pagevote"] = self.object.get_or_create_pagevote()
         return context