Add generic paginator for API

4 files changed, 105 insertions, 84 deletions

diff --git a/api/pagination.py b/api/pagination.py
new file mode 100644
index 0000000..0539ae8
--- /dev/null
+++ b/api/pagination.py
@@ -0,0 +1,45 @@
+class MastodonPaginator:
+    """
+    Paginates in the Mastodon style (max_id, min_id, etc)
+    """
+
+    def __init__(
+        self,
+        anchor_model,
+        sort_attribute: str = "created",
+        default_limit: int = 20,
+        max_limit: int = 40,
+    ):
+        self.anchor_model = anchor_model
+        self.sort_attribute = sort_attribute
+        self.default_limit = default_limit
+        self.max_limit = max_limit
+
+    def paginate(
+        self,
+        queryset,
+        min_id: str | None,
+        max_id: str | None,
+        since_id: str | None,
+        limit: int | None,
+    ):
+        if max_id:
+            anchor = self.anchor_model.objects.get(pk=max_id)
+            queryset = queryset.filter(
+                **{self.sort_attribute + "__lt": getattr(anchor, self.sort_attribute)}
+            )
+        if since_id:
+            anchor = self.anchor_model.objects.get(pk=since_id)
+            queryset = queryset.filter(
+                **{self.sort_attribute + "__gt": getattr(anchor, self.sort_attribute)}
+            )
+        if min_id:
+            # Min ID requires items _immediately_ newer than specified, so we
+            # invert the ordering to accomodate
+            anchor = self.anchor_model.objects.get(pk=min_id)
+            queryset = queryset.filter(
+                **{self.sort_attribute + "__gt": getattr(anchor, self.sort_attribute)}
+            ).order_by(self.sort_attribute)
+        else:
+            queryset = queryset.order_by("-" + self.sort_attribute)
+        return list(queryset[: min(limit or self.default_limit, self.max_limit)])
diff --git a/api/views/accounts.py b/api/views/accounts.py
index 43ec75d..4f1903b 100644
--- a/api/views/accounts.py
+++ b/api/views/accounts.py
@@ -3,6 +3,7 @@ from django.shortcuts import get_object_or_404
 from activities.models import Post, PostInteraction
 from api import schemas
 from api.decorators import identity_required
+from api.pagination import MastodonPaginator
 from api.views.base import api_router
 from users.models import Identity
 
@@ -67,7 +68,7 @@ def account_statuses(
     limit: int = 20,
 ):
     identity = get_object_or_404(Identity, pk=id)
-    posts = (
+    queryset = (
         identity.posts.not_hidden()
         .unlisted(include_replies=not exclude_replies)
         .select_related("author")
@@ -77,20 +78,16 @@ def account_statuses(
     if pinned:
         return []
     if only_media:
-        posts = posts.filter(attachments__pk__isnull=False)
+        queryset = queryset.filter(attachments__pk__isnull=False)
     if tagged:
-        posts = posts.tagged_with(tagged)
-    if max_id:
-        anchor_post = Post.objects.get(pk=max_id)
-        posts = posts.filter(created__lt=anchor_post.created)
-    if since_id:
-        anchor_post = Post.objects.get(pk=since_id)
-        posts = posts.filter(created__gt=anchor_post.created)
-    if min_id:
-        # Min ID requires LIMIT posts _immediately_ newer than specified, so we
-        # invert the ordering to accomodate
-        anchor_post = Post.objects.get(pk=min_id)
-        posts = posts.filter(created__gt=anchor_post.created).order_by("created")
-    posts = list(posts[:limit])
+        queryset = queryset.tagged_with(tagged)
+    paginator = MastodonPaginator(Post)
+    posts = paginator.paginate(
+        queryset,
+        min_id=min_id,
+        max_id=max_id,
+        since_id=since_id,
+        limit=limit,
+    )
     interactions = PostInteraction.get_post_interactions(posts, request.identity)
-    return [post.to_mastodon_json(interactions=interactions) for post in posts]
+    return [post.to_mastodon_json(interactions=interactions) for post in queryset]
diff --git a/api/views/notifications.py b/api/views/notifications.py
index 9f1f865..0b7064c 100644
--- a/api/views/notifications.py
+++ b/api/views/notifications.py
@@ -1,6 +1,7 @@
-from activities.models import Post, PostInteraction, TimelineEvent
+from activities.models import PostInteraction, TimelineEvent
 from api import schemas
 from api.decorators import identity_required
+from api.pagination import MastodonPaginator
 from api.views.base import api_router
 
 
@@ -14,8 +15,6 @@ def notifications(
     limit: int = 20,
     account_id: str | None = None,
 ):
-    if limit > 40:
-        limit = 40
     # Types/exclude_types use weird syntax so we have to handle them manually
     base_types = {
         "favourite": TimelineEvent.Types.liked,
@@ -29,7 +28,7 @@ def notifications(
         requested_types = set(base_types.keys())
     requested_types.difference_update(excluded_types)
     # Use that to pull relevant events
-    events = (
+    queryset = (
         TimelineEvent.objects.filter(
             identity=request.identity,
             type__in=[base_types[r] for r in requested_types],
@@ -37,18 +36,14 @@ def notifications(
         .order_by("-created")
         .select_related("subject_post", "subject_post__author", "subject_identity")
     )
-    if max_id:
-        anchor_post = Post.objects.get(pk=max_id)
-        events = events.filter(created__lt=anchor_post.created)
-    if since_id:
-        anchor_post = Post.objects.get(pk=since_id)
-        events = events.filter(created__gt=anchor_post.created)
-    if min_id:
-        # Min ID requires LIMIT events _immediately_ newer than specified, so we
-        # invert the ordering to accomodate
-        anchor_post = Post.objects.get(pk=min_id)
-        events = events.filter(created__gt=anchor_post.created).order_by("created")
-    events = list(events[:limit])
+    paginator = MastodonPaginator(TimelineEvent)
+    events = paginator.paginate(
+        queryset,
+        min_id=min_id,
+        max_id=max_id,
+        since_id=since_id,
+        limit=limit,
+    )
     interactions = PostInteraction.get_event_interactions(events, request.identity)
     return [
         event.to_mastodon_notification_json(interactions=interactions)
diff --git a/api/views/timelines.py b/api/views/timelines.py
index 84eed7a..8f4ac78 100644
--- a/api/views/timelines.py
+++ b/api/views/timelines.py
@@ -1,8 +1,8 @@
 from activities.models import Post, PostInteraction, TimelineEvent
-
-from .. import schemas
-from ..decorators import identity_required
-from .base import api_router
+from api import schemas
+from api.decorators import identity_required
+from api.pagination import MastodonPaginator
+from api.views.base import api_router
 
 
 @api_router.get("/v1/timelines/home", response=list[schemas.Status])
@@ -14,9 +14,8 @@ def home(
     min_id: str | None = None,
     limit: int = 20,
 ):
-    if limit > 40:
-        limit = 40
-    events = (
+    paginator = MastodonPaginator(Post)
+    queryset = (
         TimelineEvent.objects.filter(
             identity=request.identity,
             type__in=[TimelineEvent.Types.post],
@@ -25,18 +24,13 @@ def home(
         .prefetch_related("subject_post__attachments")
         .order_by("-created")
     )
-    if max_id:
-        anchor_post = Post.objects.get(pk=max_id)
-        events = events.filter(created__lt=anchor_post.created)
-    if since_id:
-        anchor_post = Post.objects.get(pk=since_id)
-        events = events.filter(created__gt=anchor_post.created)
-    if min_id:
-        # Min ID requires LIMIT events _immediately_ newer than specified, so we
-        # invert the ordering to accomodate
-        anchor_post = Post.objects.get(pk=min_id)
-        events = events.filter(created__gt=anchor_post.created).order_by("created")
-    events = list(events[:limit])
+    events = paginator.paginate(
+        queryset,
+        min_id=min_id,
+        max_id=max_id,
+        since_id=since_id,
+        limit=limit,
+    )
     interactions = PostInteraction.get_event_interactions(events, request.identity)
     return [
         event.subject_post.to_mastodon_json(interactions=interactions)
@@ -56,32 +50,26 @@ def public(
     min_id: str | None = None,
     limit: int = 20,
 ):
-    if limit > 40:
-        limit = 40
-    posts = (
+    queryset = (
         Post.objects.public()
         .select_related("author")
         .prefetch_related("attachments")
         .order_by("-created")
     )
     if local:
-        posts = posts.filter(local=True)
+        queryset = queryset.filter(local=True)
     elif remote:
-        posts = posts.filter(local=False)
+        queryset = queryset.filter(local=False)
     if only_media:
-        posts = posts.filter(attachments__id__isnull=True)
-    if max_id:
-        anchor_post = Post.objects.get(pk=max_id)
-        posts = posts.filter(created__lt=anchor_post.created)
-    if since_id:
-        anchor_post = Post.objects.get(pk=since_id)
-        posts = posts.filter(created__gt=anchor_post.created)
-    if min_id:
-        # Min ID requires LIMIT posts _immediately_ newer than specified, so we
-        # invert the ordering to accomodate
-        anchor_post = Post.objects.get(pk=min_id)
-        posts = posts.filter(created__gt=anchor_post.created).order_by("created")
-    posts = list(posts[:limit])
+        queryset = queryset.filter(attachments__id__isnull=True)
+    paginator = MastodonPaginator(Post)
+    posts = paginator.paginate(
+        queryset,
+        min_id=min_id,
+        max_id=max_id,
+        since_id=since_id,
+        limit=limit,
+    )
     interactions = PostInteraction.get_post_interactions(posts, request.identity)
     return [post.to_mastodon_json(interactions=interactions) for post in posts]
 
@@ -100,7 +88,7 @@ def hashtag(
 ):
     if limit > 40:
         limit = 40
-    posts = (
+    queryset = (
         Post.objects.public()
         .tagged_with(hashtag)
         .select_related("author")
@@ -108,21 +96,17 @@ def hashtag(
         .order_by("-created")
     )
     if local:
-        posts = posts.filter(local=True)
+        queryset = queryset.filter(local=True)
     if only_media:
-        posts = posts.filter(attachments__id__isnull=True)
-    if max_id:
-        anchor_post = Post.objects.get(pk=max_id)
-        posts = posts.filter(created__lt=anchor_post.created)
-    if since_id:
-        anchor_post = Post.objects.get(pk=since_id)
-        posts = posts.filter(created__gt=anchor_post.created)
-    if min_id:
-        # Min ID requires LIMIT posts _immediately_ newer than specified, so we
-        # invert the ordering to accomodate
-        anchor_post = Post.objects.get(pk=min_id)
-        posts = posts.filter(created__gt=anchor_post.created).order_by("created")
-    posts = list(posts[:limit])
+        queryset = queryset.filter(attachments__id__isnull=True)
+    paginator = MastodonPaginator(Post)
+    posts = paginator.paginate(
+        queryset,
+        min_id=min_id,
+        max_id=max_id,
+        since_id=since_id,
+        limit=limit,
+    )
     interactions = PostInteraction.get_post_interactions(posts, request.identity)
     return [post.to_mastodon_json(interactions=interactions) for post in posts]