feat: removable annotations on pagination

futurex_openedx_extensions/helpers/apps.py

@@ -29,4 +29,6 @@ class HelpersConfig(AppConfig):
 
     def ready(self) -> None:
         """Connect handlers to send notifications about discussions."""
+        from futurex_openedx_extensions.helpers import \
+            monkey_patches  # pylint: disable=unused-import, import-outside-toplevel
         from futurex_openedx_extensions.helpers import signals  # pylint: disable=unused-import, import-outside-toplevel

futurex_openedx_extensions/helpers/monkey_patches.py

@@ -0,0 +1,21 @@
+"""Monkey patches defined here."""
+from __future__ import annotations
+
+from typing import Any
+
+from django.db.models.query import QuerySet
+
+original_queryset_chain = QuerySet._chain  # pylint: disable=protected-access
+
+
+def customized_queryset_chain(self: Any, **kwargs: Any) -> QuerySet:
+    """Customized queryset chain method for the QuerySet class."""
+    result = original_queryset_chain(self, **kwargs)
+
+    if hasattr(self, 'removable_annotations'):
+        result.removable_annotations = self.removable_annotations.copy()
+
+    return result
+
+
+QuerySet._chain = customized_queryset_chain  # pylint: disable=protected-access

futurex_openedx_extensions/helpers/pagination.py

@@ -1,9 +1,28 @@
 """Pagination helpers and classes for the API views."""
+from django.core.paginator import Paginator
+from django.db.models.query import QuerySet
+from django.utils.functional import cached_property
 from rest_framework.pagination import PageNumberPagination
 
 
+class DefaultPaginator(Paginator):
+    """Default paginator settings for the API views."""
+    @cached_property
+    def count(self) -> int:
+        """Return the total number of objects, across all pages."""
+        if isinstance(self.object_list, QuerySet) and hasattr(self.object_list, 'removable_annotations'):
+            clone = self.object_list._chain()  # pylint: disable=protected-access
+            for key in self.object_list.removable_annotations:
+                clone.query.annotations.pop(key, None)
+            return clone.count()
+
+        return super().count
+
+
 class DefaultPagination(PageNumberPagination):
     """Default pagination settings for the API views."""
     page_size: int = 20
     page_size_query_param: str = 'page_size'
     max_page_size: int = 100
+
+    django_paginator_class = DefaultPaginator

futurex_openedx_extensions/helpers/querysets.py

@@ -18,6 +18,42 @@
 from futurex_openedx_extensions.helpers.users import get_user_by_key
 
 
+def update_removable_annotations(
+    queryset: QuerySet,
+    removable: set | List[str] | None = None,
+    not_removable: set | List[str] | None = None,
+) -> None:
+    """
+    Update the removable annotations on the given queryset.
+
+    :param queryset: QuerySet to update
+    :type queryset: QuerySet
+    :param removable: Set of annotations to add to the removable annotations
+    :type removable: set(str) | List[str] | None
+    :param not_removable: Set of annotations to remove from the removable annotations
+    :type not_removable: set(str) | List[str] | None
+    """
+    removable_annotations = queryset.removable_annotations if hasattr(queryset, 'removable_annotations') else set()
+    removable_annotations = (removable_annotations | set(removable or [])) - set(not_removable or [])
+
+    if not removable_annotations and hasattr(queryset, 'removable_annotations'):
+        delattr(queryset, 'removable_annotations')
+
+    elif removable_annotations:
+        queryset.removable_annotations = removable_annotations
+
+
+def clear_removable_annotations(queryset: QuerySet) -> None:
+    """
+    Clear the removable annotations on the given queryset.
+
+    :param queryset: QuerySet to clear the removable annotations from
+    :type queryset: QuerySet
+    """
+    if hasattr(queryset, 'removable_annotations'):
+        del queryset.removable_annotations
+
+
 def check_staff_exist_queryset(
     ref_user_id: str | Value,
     ref_org: str | Value | List | None,

tests/test_helpers/test_pagination.py

@@ -1,7 +1,10 @@
 """Tests for pagination helpers"""
+from unittest.mock import MagicMock, PropertyMock, patch
+
+from django.db.models import QuerySet
 from rest_framework.pagination import PageNumberPagination
 
-from futurex_openedx_extensions.helpers.pagination import DefaultPagination
+from futurex_openedx_extensions.helpers.pagination import DefaultPagination, DefaultPaginator
 
 
 def test_default_pagination():
@@ -10,3 +13,46 @@ def test_default_pagination():
     assert DefaultPagination.page_size == 20
     assert DefaultPagination.page_size_query_param == 'page_size'
     assert DefaultPagination.max_page_size == 100
+    assert DefaultPagination.django_paginator_class == DefaultPaginator
+
+
+@patch('futurex_openedx_extensions.helpers.pagination.Paginator.count', new_callable=PropertyMock)
+def test_count_with_query_set_and_removable_annotations(mock_super_count):
+    mock_super_count.return_value = 'should not be reached'
+
+    mock_queryset = MagicMock(spec=QuerySet)
+    mock_queryset.removable_annotations = {'annotation1'}
+    mock_queryset._chain.return_value = mock_queryset
+    mock_queryset.query.annotations = {'annotation1': None, 'annotation2': None}
+    mock_queryset.count.return_value = 5
+
+    paginator = DefaultPaginator(mock_queryset, per_page=10)
+
+    assert paginator.count == 5
+    assert 'annotation1' not in mock_queryset.query.annotations
+    mock_super_count.assert_not_called()
+
+
+@patch('futurex_openedx_extensions.helpers.pagination.Paginator.count', new_callable=PropertyMock)
+def test_count_with_query_set_no_removable_annotations(mock_super_count):
+    mock_super_count.return_value = 44
+
+    mock_queryset = MagicMock(spec=QuerySet)
+    delattr(mock_queryset, 'removable_annotations')
+    mock_queryset.count.return_value = 'should not be reached'
+    mock_queryset.query.annotations = {'annotation1': None, 'annotation2': None}
+
+    paginator = DefaultPaginator(mock_queryset, per_page=10)
+
+    assert paginator.count == mock_super_count.return_value
+    assert mock_queryset.query.annotations == {'annotation1': None, 'annotation2': None}
+
+
+@patch('futurex_openedx_extensions.helpers.pagination.Paginator.count', new_callable=PropertyMock)
+def test_count_with_not_query_set(mock_super_count):
+    mock_super_count.return_value = 44
+
+    object_list = 'not a queryset'
+    paginator = DefaultPaginator(object_list, per_page=10)
+
+    assert paginator.count == mock_super_count.return_value

tests/test_helpers/test_querysets.py

@@ -1,4 +1,6 @@
 """Tests for querysets helpers"""
+from unittest.mock import Mock
+
 import pytest
 from common.djangoapps.student.models import CourseAccessRole, UserProfile
 from django.contrib.auth import get_user_model
@@ -217,3 +219,40 @@ def test_get_permitted_learners_queryset(
 
     result = querysets.get_permitted_learners_queryset(queryset, fx_permission_info, include_staff=True)
     assert result.count() == expected_with_staff
+
+
+@pytest.mark.parametrize('original, removable, not_removable, expected_result', [
+    (None, None, None, None),
+    (None, {'f1', 'f2'}, None, {'f1', 'f2'}),
+    (None, None, {'f1', 'f2'}, None),
+    ({'f1', 'f2'}, None, None, {'f1', 'f2'}),
+    ({'f1', 'f2'}, {'f1'}, None, {'f1', 'f2'}),
+    ({'f1', 'f2'}, None, {'f1'}, {'f2'}),
+    ({'f1', 'f2'}, None, {'f1', 'f3'}, {'f2'}),
+    ({'f1', 'f2'}, {'f1'}, {'f1', 'f2'}, None),
+])
+def test_update_removable_annotations(original, removable, not_removable, expected_result):
+    """Verify update_removable_annotations function."""
+    queryset = Mock()
+    if original is None:
+        delattr(queryset, 'removable_annotations')
+    else:
+        queryset.removable_annotations = original
+
+    querysets.update_removable_annotations(queryset, removable, not_removable)
+
+    if expected_result is None:
+        assert not hasattr(queryset, 'removable_annotations')
+    else:
+        assert queryset.removable_annotations == expected_result
+
+
+def test_clear_removable_annotations():
+    """Verify clear_removable_annotations function."""
+    queryset = Mock(removable_annotations={'f1', 'f2'})
+    querysets.clear_removable_annotations(queryset)
+
+    assert not hasattr(queryset, 'removable_annotations')
+
+    querysets.clear_removable_annotations(queryset)
+    assert not hasattr(queryset, 'removable_annotations')