squash!: Generalize user assignments filters

Author: rodmgwgu

openedx_authz/api/data.py

@@ -1183,3 +1183,9 @@ def __repr__(self):
         """Developer friendly string representation of the role assignment."""
         role_keys = ", ".join(role.namespaced_key for role in self.roles)
         return f"{self.subject.namespaced_key} => [{role_keys}] @ {self.scope.namespaced_key}"
+
+
+@define
+class UserAssignments:
+    user: "User"
+    assignments: list[RoleAssignmentData]

openedx_authz/rest_api/data.py

@@ -50,3 +50,10 @@ class RoleOperationError(BaseEnum):
     USER_DOES_NOT_HAVE_ROLE = "user_does_not_have_role"
     ROLE_ASSIGNMENT_ERROR = "role_assignment_error"
     ROLE_REMOVAL_ERROR = "role_removal_error"
+
+
+class UserAssignmentsFilter(BaseEnum):
+    """Enum for the filters that can be applied over UserAssignments."""
+
+    SCOPES = "scopes"
+    ORGS = "orgs"

openedx_authz/rest_api/utils.py

@@ -1,7 +1,6 @@
 """Utility functions for the Open edX AuthZ REST API."""
 
 import logging
-from dataclasses import dataclass
 
 from django.contrib.auth import get_user_model
 
@@ -13,10 +12,11 @@
     OrgCourseOverviewGlobData,
     RoleAssignmentData,
     ScopeData,
+    UserAssignments,
 )
 from openedx_authz.api.users import is_user_allowed
 from openedx_authz.constants.permissions import COURSES_VIEW_COURSE, VIEW_LIBRARY
-from openedx_authz.rest_api.data import SearchField, SortField, SortOrder
+from openedx_authz.rest_api.data import SearchField, SortField, SortOrder, UserAssignmentsFilter
 
 logger = logging.getLogger(__name__)
 
@@ -131,12 +131,6 @@ def filter_users(users: list[dict], search: str | None, roles: list[str] | None)
     return filtered_users
 
 
-@dataclass
-class UserAssignments:
-    user: "User"
-    assignments: list[RoleAssignmentData]
-
-
 def get_user_assignment_map(role_assignments: list[RoleAssignmentData]) -> list[UserAssignments]:
     """
     Group role assignments by user
@@ -175,3 +169,33 @@ def filter_allowed_assignments(user: "User", assignments: list[RoleAssignmentDat
             allowed_assignments.append(assignment)
 
     return allowed_assignments
+
+
+def filter_user_assignments(
+    users_with_assignments: list[UserAssignments],
+    by: UserAssignmentsFilter,
+    values: list[str],
+) -> list[UserAssignments]:
+    """
+    Filter user assignments by orgs or scopes.
+    """
+    if not values:
+        return users_with_assignments
+
+    def _get_value_to_filter(assignment: RoleAssignmentData) -> str:
+        if by == UserAssignmentsFilter.SCOPES:
+            return assignment.scope.external_key
+        elif by == UserAssignmentsFilter.ORGS:
+            return assignment.scope.org
+        else:
+            raise ValueError(f"Invalid filter: '{by}'. Must be one of {UserAssignmentsFilter.values()}")
+
+    filtered_users: list[UserAssignments] = []
+    for uwa in users_with_assignments:
+        if any(_get_value_to_filter(a) in values for a in uwa.assignments):
+            # Also filter assignments to reflect the correct number of assignments
+            filtered_assignments = [a for a in uwa.assignments if _get_value_to_filter(a) in values]
+            filtered_users.append(UserAssignments(user=uwa.user, assignments=filtered_assignments))
+    users_with_assignments = filtered_users
+
+    return filtered_users

openedx_authz/rest_api/v1/views.py

@@ -16,11 +16,11 @@
 
 from openedx_authz import api
 from openedx_authz.constants import permissions
-from openedx_authz.rest_api.data import RoleOperationError, RoleOperationStatus, SortField
+from openedx_authz.rest_api.data import RoleOperationError, RoleOperationStatus, SortField, UserAssignmentsFilter
 from openedx_authz.rest_api.decorators import authz_permissions, view_auth_classes
 from openedx_authz.rest_api.utils import (
-    UserAssignments,
     filter_allowed_assignments,
+    filter_user_assignments,
     filter_users,
     get_generic_scope,
     get_user_assignment_map,
@@ -500,25 +500,19 @@ def get(self, request: HttpRequest) -> Response:
         order = query_params.get("order")
         sort_by = query_params.get("sort_by", SortField.USERNAME)
 
-        if scopes:
-            # Filter by scopes
-            filtered_users: list[UserAssignments] = []
-            for uwa in users_with_assignments:
-                if any(a.scope.external_key in scopes for a in uwa.assignments):
-                    # Also filter assignments to reflect the correct number of assignments
-                    filtered_assignments = [a for a in uwa.assignments if a.scope.external_key in scopes]
-                    filtered_users.append(UserAssignments(user=uwa.user, assignments=filtered_assignments))
-            users_with_assignments = filtered_users
-
-        if orgs:
-            # Filter by orgs
-            filtered_users: list[UserAssignments] = []
-            for uwa in users_with_assignments:
-                if any(a.scope.org in orgs for a in uwa.assignments):
-                    # Also filter assignments to reflect the correct number of assignments
-                    filtered_assignments = [a for a in uwa.assignments if a.scope.org in orgs]
-                    filtered_users.append(UserAssignments(user=uwa.user, assignments=filtered_assignments))
-            users_with_assignments = filtered_users
+        # Filter by scopes
+        users_with_assignments = filter_user_assignments(
+            users_with_assignments=users_with_assignments,
+            by=UserAssignmentsFilter.SCOPES,
+            values=scopes,
+        )
+
+        # Filter by orgs
+        users_with_assignments = filter_user_assignments(
+            users_with_assignments=users_with_assignments,
+            by=UserAssignmentsFilter.ORGS,
+            values=orgs,
+        )
 
         team_members = TeamMemberSerializer(users_with_assignments, many=True)
         # Search