Merge pull request #800 from superannotateai/FRIDAY-3995

added contributor categories functions

Author: nareksa

docs/source/api_reference/api_team.rst

@@ -15,3 +15,5 @@ Team
 .. automethod:: superannotate.SAClient.resume_user_activity
 .. automethod:: superannotate.SAClient.get_user_scores
 .. automethod:: superannotate.SAClient.set_user_scores
+.. automethod:: superannotate.SAClient.set_contributors_categories
+.. automethod:: superannotate.SAClient.remove_contributors_categories

src/superannotate/lib/app/interface/sdk_interface.py

@@ -474,7 +474,7 @@ def list_users(
         self,
         *,
         project: Union[int, str] = None,
-        include: List[Literal["custom_fields"]] = None,
+        include: List[Literal["custom_fields", "categories"]] = None,
         **filters,
     ):
         """
@@ -488,7 +488,10 @@ def list_users(
 
             Possible values are
 
-            - "custom_fields":  Includes custom fields and scores assigned to each user.
+            - "custom_fields": Includes custom fields and scores assigned to each user.
+            - "categories": Includes a list of categories assigned to each project contributor.
+              Note: 'project' parameter must be specified when including 'categories'.
+        :type include: list of str, optional
 
         :param filters: Specifies filtering criteria, with all conditions combined using logical AND.
 
@@ -860,6 +863,103 @@ def set_user_scores(
         )
         logger.info("Scores successfully set.")
 
+    def set_contributors_categories(
+        self,
+        project: Union[NotEmptyStr, int],
+        contributors: List[Union[int, str]],
+        categories: Union[List[str], Literal["*"]],
+    ):
+        """
+        Assign one or more categories to a contributor with an assignable role (Annotator, QA or custom role)
+        in a Multimodal project. Project Admins are not eligible for category assignments. "*" in the category
+        list will match all categories defined in the project.
+
+
+        :param project: The name or ID of the project.
+        :type project: Union[NotEmptyStr, int]
+
+        :param contributors: A list of emails or IDs of the contributor.
+        :type contributors: List[Union[int, str]]
+
+        :param categories:  A list of category names to assign. Accepts "*" to indicate all available categories in the project.
+        :type categories: Union[List[str], Literal["*"]]
+
+        Request Example:
+        ::
+
+            client.set_contributor_categories(
+                project="product-review-mm",
+                contributors=["test@superannotate.com","contributor@superannotate.com"],
+                categories=["Shoes", "T-Shirt"]
+            )
+
+            client.set_contributor_categories(
+                project="product-review-mm",
+                contributors=["test@superannotate.com","contributor@superannotate.com"]
+                categories="*"
+            )
+        """
+        project = (
+            self.controller.get_project_by_id(project).data
+            if isinstance(project, int)
+            else self.controller.get_project(project)
+        )
+        self.controller.check_multimodal_project_categorization(project)
+
+        self.controller.work_management.set_remove_contributor_categories(
+            project=project,
+            contributors=contributors,
+            categories=categories,
+            operation="set",
+        )
+
+    def remove_contributors_categories(
+        self,
+        project: Union[NotEmptyStr, int],
+        contributors: List[Union[int, str]],
+        categories: Union[List[str], Literal["*"]],
+    ):
+        """
+        Remove one or more categories for a contributor. "*" in the category list will match all categories defined in the project.
+
+        :param project: The name or ID of the project.
+        :type project: Union[NotEmptyStr, int]
+
+        :param contributors: A list of emails or IDs of the contributor.
+        :type contributors: List[Union[int, str]]
+
+        :param categories:  A list of category names to remove. Accepts "*" to indicate all available categories in the project.
+        :type categories: Union[List[str], Literal["*"]]
+
+        Request Example:
+        ::
+
+            client.remove_contributor_categories(
+                project="product-review-mm",
+                contributors=["test@superannotate.com","contributor@superannotate.com"],
+                categories=["Shoes", "T-Shirt", "Jeans"]
+            )
+
+            client.remove_contributor_categories(
+                project="product-review-mm",
+                contributors=["test@superannotate.com","contributor@superannotate.com"]
+                categories="*"
+            )
+        """
+        project = (
+            self.controller.get_project_by_id(project).data
+            if isinstance(project, int)
+            else self.controller.get_project(project)
+        )
+        self.controller.check_multimodal_project_categorization(project)
+
+        self.controller.work_management.set_remove_contributor_categories(
+            project=project,
+            contributors=contributors,
+            categories=categories,
+            operation="remove",
+        )
+
     def get_component_config(self, project: Union[NotEmptyStr, int], component_id: str):
         """
         Retrieves the configuration for a given project and component ID.
@@ -1320,6 +1420,7 @@ def remove_categories(
         )
         self.controller.check_multimodal_project_categorization(project)
 
+        categories_to_remove = None
         query = EmptyQuery()
         if categories == "*":
             query &= Filter("id", [0], OperatorEnum.GT)
@@ -1335,14 +1436,13 @@ def remove_categories(
         else:
             raise AppException("Categories should be a list of strings or '*'.")
 
-        response = (
-            self.controller.service_provider.work_management.remove_project_categories(
+        if categories_to_remove:
+            response = self.controller.service_provider.work_management.remove_project_categories(
                 project_id=project.id, query=query
             )
-        )
-        logger.info(
-            f"{len(response.data)} categories successfully removed from the project."
-        )
+            logger.info(
+                f"{len(response.data)} categories successfully removed from the project."
+            )
 
     def create_folder(self, project: NotEmptyStr, folder_name: NotEmptyStr):
         """

src/superannotate/lib/core/entities/work_managament.py

@@ -129,11 +129,12 @@ def json(self, **kwargs):
 class WMProjectUserEntity(TimedBaseModel):
     id: Optional[int]
     team_id: Optional[int]
-    role: int
+    role: Optional[int]
     email: Optional[str]
     state: Optional[WMUserStateEnum]
     custom_fields: Optional[dict] = Field(dict(), alias="customField")
     permissions: Optional[dict]
+    categories: Optional[list[dict]]
 
     class Config:
         extra = Extra.ignore

src/superannotate/lib/core/serviceproviders.py

@@ -228,6 +228,17 @@ def create_score(
     def delete_score(self, score_id: int) -> ServiceResponse:
         raise NotImplementedError
 
+    @abstractmethod
+    def set_remove_contributor_categories(
+        self,
+        project_id: int,
+        contributor_ids: List[int],
+        category_ids: List[int],
+        operation: Literal["set", "remove"],
+        chunk_size=100,
+    ) -> list[dict]:
+        raise NotImplementedError
+
 
 class BaseProjectService(SuperannotateServiceProvider):
     @abstractmethod

src/superannotate/lib/infrastructure/controller.py

@@ -68,6 +68,9 @@
 from typing_extensions import Unpack
 
 
+logger = logging.getLogger("sa")
+
+
 def build_condition(**kwargs) -> Condition:
     condition = Condition.get_empty_condition()
     if any(kwargs.values()):
@@ -177,7 +180,10 @@ def set_custom_field_value(
         )
 
     def list_users(
-        self, include: List[Literal["custom_fields"]] = None, project=None, **filters
+        self,
+        include: List[Literal["custom_fields", "categories"]] = None,
+        project=None,
+        **filters,
     ):
         context = {"team_id": self.service_provider.client.team_id}
         if project:
@@ -205,6 +211,10 @@ def list_users(
             ]
         )
         query = chain.handle(filters, EmptyQuery())
+
+        if project and include and "categories" in include:
+            query &= Join("categories")
+
         if include and "custom_fields" in include:
             response = self.service_provider.work_management.list_users(
                 query,
@@ -401,6 +411,76 @@ def set_user_scores(
             res.res_error = "Please provide valid score values."
         res.raise_for_status()
 
+    def set_remove_contributor_categories(
+        self,
+        project: ProjectEntity,
+        contributors: List[Union[int, str]],
+        categories: Union[List[str], Literal["*"]],
+        operation: Literal["set", "remove"],
+    ):
+        if categories and contributors:
+            all_categories = (
+                self.service_provider.work_management.list_project_categories(
+                    project_id=project.id, entity=ProjectCategoryEntity  # noqa
+                ).data
+            )
+            if categories == "*":
+                category_ids = [c.id for c in all_categories]
+            else:
+                categories = [c.lower() for c in categories]
+                category_ids = [
+                    c.id for c in all_categories if c.name.lower() in categories
+                ]
+
+            if isinstance(contributors[0], str):
+                project_contributors = self.list_users(
+                    project=project, email__in=contributors
+                )
+            elif isinstance(contributors[0], int):
+                project_contributors = self.list_users(
+                    project=project, id__in=contributors
+                )
+            else:
+                raise AppException("Contributors not found.")
+
+            if len(project_contributors) < len(contributors):
+                raise AppException("Contributors not found.")
+
+            contributor_ids = [
+                c.id
+                for c in project_contributors
+                if c.role != 3  # exclude Project Admins
+            ]
+
+            if category_ids and contributor_ids:
+                response = self.service_provider.work_management.set_remove_contributor_categories(
+                    project_id=project.id,
+                    contributor_ids=contributor_ids,
+                    category_ids=category_ids,
+                    operation=operation,
+                )
+
+                success_processed = 0
+                for contributor in response:
+                    contributor_category_ids = [
+                        category["id"] for category in contributor["categories"]
+                    ]
+                    if operation == "set":
+                        if set(category_ids).issubset(contributor_category_ids):
+                            success_processed += len(category_ids)
+                    else:
+                        if not set(category_ids).intersection(contributor_category_ids):
+                            success_processed += len(category_ids)
+
+                if success_processed / len(contributor_ids) == len(category_ids):
+                    action_for_log = (
+                        "added to" if operation == "set" else "removed from"
+                    )
+                    logger.info(
+                        f"{len(category_ids)} categories successfully {action_for_log} "
+                        f"{len(contributor_ids)} contributors."
+                    )
+
 
 class ProjectManager(BaseManager):
     def __init__(self, service_provider: ServiceProvider, team: TeamEntity):

src/superannotate/lib/infrastructure/services/work_management.py

@@ -14,6 +14,7 @@
 from lib.core.entities.work_managament import WMUserEntity
 from lib.core.enums import CustomFieldEntityEnum
 from lib.core.exceptions import AppException
+from lib.core.jsx_conditions import EmptyQuery
 from lib.core.jsx_conditions import Filter
 from lib.core.jsx_conditions import OperatorEnum
 from lib.core.jsx_conditions import Query
@@ -71,6 +72,7 @@ class WorkManagementService(BaseWorkManagementService):
     URL_SEARCH_PROJECT_USERS = "projectusers/search"
     URL_SEARCH_PROJECTS = "projects/search"
     URL_RESUME_PAUSE_USER = "teams/editprojectsusers"
+    URL_CONTRIBUTORS_CATEGORIES = "customentities/edit"
 
     @staticmethod
     def _generate_context(**kwargs):
@@ -475,3 +477,46 @@ def delete_score(self, score_id: int) -> ServiceResponse:
                 ),
             },
         )
+
+    def set_remove_contributor_categories(
+        self,
+        project_id: int,
+        contributor_ids: List[int],
+        category_ids: List[int],
+        operation: Literal["set", "remove"],
+        chunk_size=100,
+    ) -> List[dict]:
+        params = {
+            "entity": "Contributor",
+            "parentEntity": "Project",
+        }
+        if operation == "set":
+            params["action"] = "addcontributorcategory"
+        else:
+            params["action"] = "removecontributorcategory"
+
+        from lib.infrastructure.utils import divide_to_chunks
+
+        success_contributors = []
+
+        for chunk in divide_to_chunks(contributor_ids, chunk_size):
+            body_query = EmptyQuery()
+            body_query &= Filter("id", chunk, OperatorEnum.IN)
+            response = self.client.request(
+                url=self.URL_CONTRIBUTORS_CATEGORIES,
+                method="post",
+                params=params,
+                data={
+                    "query": body_query.body_builder(),
+                    "body": {"categories": [{"id": i} for i in category_ids]},
+                },
+                headers={
+                    "x-sa-entity-context": self._generate_context(
+                        team_id=self.client.team_id, project_id=project_id
+                    ),
+                },
+            )
+            response.raise_for_status()
+            success_contributors.extend(response.data["data"])
+
+        return success_contributors