filter by user type

Author: SKairinos

codeforlife/user/filters/user.py

@@ -12,7 +12,12 @@
 )
 
 from ...filters import FilterSet  # isort: skip
-from ..models import User  # isort: skip
+from ..models import (  # isort: skip
+    User,
+    TeacherUser,
+    StudentUser,
+    IndependentUser,
+)
 
 
 # pylint: disable-next=missing-class-docstring
@@ -22,14 +27,22 @@ class UserFilterSet(FilterSet):
         "exact",
     )
 
-    _id = filters.NumberFilter(method="_id_method")
-    _id_method = FilterSet.make_exclude_field_list_method("id")
+    _id = filters.NumberFilter(method="_id__method")
+    _id__method = FilterSet.make_exclude_field_list_method("id")
 
-    name = filters.CharFilter(method="name_method")
+    name = filters.CharFilter(method="name__method")
 
-    only_teachers = filters.BooleanFilter(method="only_teachers__method")
+    type = filters.ChoiceFilter(
+        choices=[
+            ("teacher", "teacher"),
+            ("student", "student"),
+            ("independent", "independent"),
+            ("indy", "independent"),
+        ],
+        method="type__method",
+    )
 
-    def name_method(
+    def name__method(
         self: FilterSet, queryset: QuerySet[User], name: str, *args
     ):
         """Get all first names and last names that contain a substring."""
@@ -45,16 +58,19 @@ def name_method(
             | Q(last_name__icontains=last_name)
         )
 
-    def only_teachers__method(
-        self: FilterSet, queryset: QuerySet[User], _: str, value: bool
+    def type__method(
+        self: FilterSet,
+        queryset: QuerySet[User],
+        _: str,
+        value: t.Literal["teacher", "student", "independent"],
     ):
-        """Get only teacher-users."""
-        return (
-            queryset.filter(new_teacher__isnull=False, new_student__isnull=True)
-            if value
-            else queryset
-        )
+        """Get users of a specific type."""
+        if value == "teacher":
+            return TeacherUser.objects.filter_users(queryset)
+        if value == "student":
+            return StudentUser.objects.filter_users(queryset)
+        return IndependentUser.objects.filter_users(queryset)
 
     class Meta:
         model = User
-        fields = ["students_in_class", "only_teachers", "_id", "name"]
+        fields = ["students_in_class", "type", "_id", "name"]

codeforlife/user/models/user.py

@@ -11,7 +11,7 @@
 
 # pylint: disable-next=imported-auth-user
 from django.contrib.auth.models import User as _User
-from django.contrib.auth.models import UserManager
+from django.contrib.auth.models import UserManager as _UserManager
 from django.db.models import F
 from django.db.models.query import QuerySet
 from django.utils.crypto import get_random_string
@@ -157,13 +157,28 @@ def anonymize(self):
 AnyUser = t.TypeVar("AnyUser", bound=User)
 
 
-# pylint: disable-next=missing-class-docstring,too-few-public-methods
-class ContactableUserManager(UserManager[AnyUser], t.Generic[AnyUser]):
+# pylint: disable-next=missing-class-docstring
+class UserManager(_UserManager[AnyUser], t.Generic[AnyUser]):
+    def filter_users(self, queryset: QuerySet[User]):
+        """Filter the users to the specific type.
+
+        Args:
+            queryset: The queryset of users to filter.
+
+        Returns:
+            A subset of the queryset of users.
+        """
+        return queryset
+
     # pylint: disable-next=missing-function-docstring
     def get_queryset(self):
-        return (
-            super().get_queryset().exclude(email__isnull=True).exclude(email="")
-        )
+        return self.filter_users(super().get_queryset())
+
+
+# pylint: disable-next=missing-class-docstring,too-few-public-methods
+class ContactableUserManager(UserManager[AnyUser], t.Generic[AnyUser]):
+    def filter_users(self, queryset: QuerySet[User]):
+        return queryset.exclude(email__isnull=True).exclude(email="")
 
 
 class ContactableUser(User):
@@ -238,15 +253,16 @@ def create_user(  # type: ignore[override]
 
         return user
 
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
+    def filter_users(self, queryset: QuerySet[User]):
         return (
             super()
-            .get_queryset()
+            .filter_users(queryset)
             .filter(new_teacher__isnull=False, new_student__isnull=True)
-            .prefetch_related("new_teacher")
         )
 
+    def get_queryset(self):
+        return super().get_queryset().prefetch_related("new_teacher")
+
 
 class TeacherUser(ContactableUser):
     """A user that is a teacher."""
@@ -287,9 +303,12 @@ def create_user(  # type: ignore[override]
             **extra_fields,
         )
 
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
-        return super().get_queryset().filter(new_teacher__school__isnull=False)
+    def filter_users(self, queryset: QuerySet[User]):
+        return (
+            super()
+            .filter_users(queryset)
+            .filter(new_teacher__school__isnull=False)
+        )
 
 
 # pylint: disable-next=too-many-ancestors
@@ -317,9 +336,8 @@ def teacher(self):
 class AdminSchoolTeacherUserManager(
     SchoolTeacherUserManager["AdminSchoolTeacherUser"]
 ):
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
-        return super().get_queryset().filter(new_teacher__is_admin=True)
+    def filter_users(self, queryset: QuerySet[User]):
+        return super().filter_users(queryset).filter(new_teacher__is_admin=True)
 
 
 # pylint: disable-next=too-many-ancestors
@@ -347,9 +365,10 @@ def teacher(self):
 class NonAdminSchoolTeacherUserManager(
     SchoolTeacherUserManager["NonAdminSchoolTeacherUser"]
 ):
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
-        return super().get_queryset().filter(new_teacher__is_admin=False)
+    def filter_users(self, queryset: QuerySet[User]):
+        return (
+            super().filter_users(queryset).filter(new_teacher__is_admin=False)
+        )
 
 
 # pylint: disable-next=too-many-ancestors
@@ -377,9 +396,12 @@ def teacher(self):
 
 # pylint: disable-next=missing-class-docstring,too-few-public-methods
 class NonSchoolTeacherUserManager(TeacherUserManager["NonSchoolTeacherUser"]):
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
-        return super().get_queryset().filter(new_teacher__school__isnull=True)
+    def filter_users(self, queryset: QuerySet[User]):
+        return (
+            super()
+            .filter_users(queryset)
+            .filter(new_teacher__school__isnull=True)
+        )
 
 
 # pylint: disable-next=too-many-ancestors
@@ -440,20 +462,17 @@ def create_user(  # type: ignore[override]
 
         return user
 
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
-        return (
-            super()
-            .get_queryset()
-            .filter(
-                new_teacher__isnull=True,
-                new_student__isnull=False,
-                # TODO: remove in new model
-                new_student__class_field__isnull=False,
-            )
-            .prefetch_related("new_student")
+    def filter_users(self, queryset: QuerySet[User]):
+        return queryset.filter(
+            new_teacher__isnull=True,
+            new_student__isnull=False,
+            # TODO: remove in new model
+            new_student__class_field__isnull=False,
         )
 
+    def get_queryset(self):
+        return super().get_queryset().prefetch_related("new_student")
+
 
 class StudentUser(User):
     """A user that is a student."""
@@ -509,20 +528,21 @@ def set_password(self, raw_password: t.Optional[str] = None):
 
 # pylint: disable-next=missing-class-docstring,too-few-public-methods
 class IndependentUserManager(ContactableUserManager["IndependentUser"]):
-    # pylint: disable-next=missing-function-docstring
-    def get_queryset(self):
-        # TODO: student__isnull=True in new model
+    def filter_users(self, queryset: QuerySet[User]):
         return (
             super()
-            .get_queryset()
+            .filter_users(queryset)
             .filter(
                 new_teacher__isnull=True,
+                # TODO: student__isnull=True in new model
                 new_student__isnull=False,
                 new_student__class_field__isnull=True,
             )
-            .prefetch_related("new_student")
         )
 
+    def get_queryset(self):
+        return super().get_queryset().prefetch_related("new_student")
+
     def create_user(  # type: ignore[override]
         self,
         first_name: str,