Moves privacy filter to model class

bookwyrm/models/base_model.py

@@ -2,8 +2,9 @@
 import base64
 from Crypto import Random
 
-from django.core.exceptions import PermissionDenied
+from django.core.exceptions import PermissionDenied, FieldError
 from django.db import models
+from django.db.models import Q
 from django.dispatch import receiver
 from django.http import Http404
 from django.utils.translation import gettext_lazy as _
@@ -105,6 +106,76 @@ class BookWyrmModel(models.Model):
 
         raise PermissionDenied()
 
+    @classmethod
+    def privacy_filter(
+        cls, viewer, privacy_levels=None, following_only=False, **filters
+    ):
+        """filter objects that have "user" and "privacy" fields"""
+        queryset = cls.objects
+        if hasattr(queryset, "select_subclasses"):
+            queryset = queryset.select_subclasses()
+
+        privacy_levels = privacy_levels or ["public", "unlisted", "followers", "direct"]
+        # if there'd a deleted field, exclude deleted items
+        try:
+            queryset = queryset.filter(deleted=False)
+        except FieldError:
+            pass
+
+        # exclude blocks from both directions
+        if not viewer.is_anonymous:
+            queryset = queryset.exclude(
+                Q(user__blocked_by=viewer) | Q(user__blocks=viewer)
+            )
+
+        # you can't see followers only or direct messages if you're not logged in
+        if viewer.is_anonymous:
+            privacy_levels = [
+                p for p in privacy_levels if not p in ["followers", "direct"]
+            ]
+
+        # filter to only privided privacy levels
+        queryset = queryset.filter(privacy__in=privacy_levels)
+
+        # only include statuses the user follows
+        if following_only:
+            queryset = queryset.exclude(
+                ~Q(  # remove everythign except
+                    Q(user__followers=viewer)
+                    | Q(user=viewer)  # user following
+                    | Q(mention_users=viewer)  # is self  # mentions user
+                ),
+            )
+        # exclude followers-only statuses the user doesn't follow
+        elif "followers" in privacy_levels:
+            queryset = cls.followers_filter(queryset, viewer)
+
+        # exclude direct messages not intended for the user
+        if "direct" in privacy_levels:
+            queryset = cls.direct_filter(queryset, viewer)
+
+        return queryset
+
+    @classmethod
+    def followers_filter(cls, queryset, viewer):
+        """Override-able filter for "followers" privacy level"""
+        return queryset.exclude(
+            ~Q(  # user isn't following and it isn't their own status
+                Q(user__followers=viewer) | Q(user=viewer)
+            ),
+            privacy="followers",  # and the status is followers only
+        )
+
+    @classmethod
+    def direct_filter(cls, queryset, viewer):
+        """Override-able filter for "direct" privacy level"""
+        try:
+            return queryset.exclude(
+                ~Q(Q(user=viewer) | Q(mention_users=viewer)), privacy="direct"
+            )
+        except FieldError:
+            return queryset.exclude(~Q(user=viewer), privacy="direct")
+
 
 @receiver(models.signals.post_save)
 # pylint: disable=unused-argument

bookwyrm/views/helpers.py

@@ -6,8 +6,6 @@ import dateutil.tz
 from dateutil.parser import ParserError
 
 from requests import HTTPError
-from django.core.exceptions import FieldError
-from django.db.models import Q
 from django.http import Http404
 
 from bookwyrm import activitypub, models
@@ -50,56 +48,6 @@ def is_bookwyrm_request(request):
     return True
 
 
-def privacy_filter(viewer, queryset, privacy_levels=None, following_only=False):
-    """filter objects that have "user" and "privacy" fields"""
-    privacy_levels = privacy_levels or ["public", "unlisted", "followers", "direct"]
-    # if there'd a deleted field, exclude deleted items
-    try:
-        queryset = queryset.filter(deleted=False)
-    except FieldError:
-        pass
-
-    # exclude blocks from both directions
-    if not viewer.is_anonymous:
-        queryset = queryset.exclude(Q(user__blocked_by=viewer) | Q(user__blocks=viewer))
-
-    # you can't see followers only or direct messages if you're not logged in
-    if viewer.is_anonymous:
-        privacy_levels = [p for p in privacy_levels if not p in ["followers", "direct"]]
-
-    # filter to only privided privacy levels
-    queryset = queryset.filter(privacy__in=privacy_levels)
-
-    # only include statuses the user follows
-    if following_only:
-        queryset = queryset.exclude(
-            ~Q(  # remove everythign except
-                Q(user__followers=viewer)
-                | Q(user=viewer)  # user following
-                | Q(mention_users=viewer)  # is self  # mentions user
-            ),
-        )
-    # exclude followers-only statuses the user doesn't follow
-    elif "followers" in privacy_levels:
-        queryset = queryset.exclude(
-            ~Q(  # user isn't following and it isn't their own status
-                Q(user__followers=viewer) | Q(user=viewer)
-            ),
-            privacy="followers",  # and the status is followers only
-        )
-
-    # exclude direct messages not intended for the user
-    if "direct" in privacy_levels:
-        try:
-            queryset = queryset.exclude(
-                ~Q(Q(user=viewer) | Q(mention_users=viewer)), privacy="direct"
-            )
-        except FieldError:
-            queryset = queryset.exclude(~Q(user=viewer), privacy="direct")
-
-    return queryset
-
-
 def handle_remote_webfinger(query):
     """webfingerin' other servers"""
     user = None