save

Author: Adrian Chang

libs/labelbox/src/labelbox/schema/user_group.py

@@ -1,12 +1,16 @@
 from enum import Enum
-from typing import Set, List, Union, Iterator, Optional
+from typing import Set, Iterator
+from collections import defaultdict
 
 from labelbox import Client
 from labelbox.exceptions import ResourceCreationError
 from labelbox.pydantic_compat import BaseModel
 from labelbox.schema.user import User
 from labelbox.schema.project import Project
 from labelbox.exceptions import UnprocessableEntityError, InvalidQueryError
+from labelbox.schema.queue_mode import QueueMode
+from labelbox.schema.ontology_kind import EditorTaskType
+from labelbox.schema.media_type import MediaType
 
 
 class UserGroupColor(Enum):
@@ -34,55 +38,6 @@ class UserGroupColor(Enum):
     YELLOW = "E7BF00"
     GRAY = "B8C4D3"
 
-
-class UserGroupUser(BaseModel):
-    """
-    Represents a user in a group.
-
-    Attributes:
-        id (str): The ID of the user.
-        email (str): The email of the user.
-    """
-    id: str
-    email: str
-
-    def __hash__(self):
-        return hash((self.id))
-
-    def __eq__(self, other):
-        if not isinstance(other, UserGroupUser):
-            return False
-        return self.id == other.id
-
-
-class UserGroupProject(BaseModel):
-    """
-    Represents a project in a group.
-
-    Attributes:
-        id (str): The ID of the project.
-        name (str): The name of the project.
-    """
-    id: str
-    name: str
-
-    def __hash__(self):
-        return hash((self.id))
-
-    def __eq__(self, other):
-        """
-        Check if this GroupProject object is equal to another GroupProject object.
-
-        Args:
-            other (GroupProject): The other GroupProject object to compare with.
-
-        Returns:
-            bool: True if the two GroupProject objects are equal, False otherwise.
-        """
-        if not isinstance(other, UserGroupProject):
-            return False
-        return self.id == other.id
-
 
 class UserGroup(BaseModel):
     """
@@ -107,8 +62,8 @@ class UserGroup(BaseModel):
     id: str
     name: str
     color: UserGroupColor
-    users: Set[UserGroupUser]
-    projects: Set[UserGroupProject]
+    users: Set[User]
+    projects: Set[Project]
     client: Client
 
     class Config:
@@ -182,14 +137,8 @@ def get(self) -> "UserGroup":
             raise InvalidQueryError("Failed to fetch group")
         self.name = result["userGroup"]["name"]
         self.color = UserGroupColor(result["userGroup"]["color"])
-        self.projects = {
-            UserGroupProject(id=project["id"], name=project["name"])
-            for project in result["userGroup"]["projects"]["nodes"]
-        }
-        self.users = {
-            UserGroupUser(id=member["id"], email=member["email"])
-            for member in result["userGroup"]["members"]["nodes"]
-        }
+        self.projects = self._get_projects_set(result["userGroup"]["projects"]["nodes"])
+        self.users = self._get_users_set(result["userGroup"]["members"]["nodes"])
         return self
 
     def update(self) -> "UserGroup":
@@ -236,10 +185,10 @@ def update(self) -> "UserGroup":
             "color":
                 self.color.value,
             "projectIds": [
-                project.id for project in self.projects
+                project.uid for project in self.projects
             ],
             "userIds": [
-                user.id for user in self.users
+                user.uid for user in self.users
             ]
         }
         result = self.client.execute(query, params)
@@ -298,10 +247,10 @@ def create(self) -> "UserGroup":
             "color":
                 self.color.value,
             "projectIds": [
-                project.id for project in self.projects
+                project.uid for project in self.projects
             ],
             "userIds": [
-                user.id for user in self.users
+                user.uid for user in self.users
             ]
         }
         result = self.client.execute(query, params)
@@ -387,16 +336,49 @@ def get_user_groups(self) -> Iterator["UserGroup"]:
                 userGroup.id = group["id"]
                 userGroup.name = group["name"]
                 userGroup.color = UserGroupColor(group["color"])
-                userGroup.users = {
-                    UserGroupUser(id=member["id"], email=member["email"])
-                    for member in group["members"]["nodes"]
-                }
-                userGroup.projects = {
-                    UserGroupProject(id=project["id"], name=project["name"])
-                    for project in group["projects"]["nodes"]
-                }
+                userGroup.users = self._get_users_set(group["members"]["nodes"])
+                userGroup.projects = self._get_projects_set(group["projects"]["nodes"])
                 yield userGroup
             nextCursor = userGroups["nextCursor"]
             # this doesn't seem to be implemented right now to return a value other than null from the api
             if not nextCursor:
                 break
+
+    def _get_users_set(self, user_nodes):
+        """
+        Retrieves a set of User objects from the given user nodes.
+
+        Args:
+            user_nodes (list): A list of user nodes containing user information.
+
+        Returns:
+            set: A set of User objects.
+        """
+        users = set()
+        for user in user_nodes:
+            user_values = defaultdict(lambda: None)
+            user_values["id"] = user["id"]
+            user_values["email"] = user["email"]
+            users.add(User(self.client, user_values))
+        return users
+
+    def _get_projects_set(self, project_nodes):
+        """
+        Retrieves a set of projects based on the given project nodes.
+
+        Args:
+            project_nodes (list): A list of project nodes.
+
+        Returns:
+            set: A set of Project objects.
+        """
+        projects = set()
+        for project in project_nodes:
+            project_values = defaultdict(lambda: None)
+            project_values["id"] = project["id"]
+            project_values["name"] = project["name"]
+            project_values["queueMode"] = QueueMode.Batch.value
+            project_values["editorTaskType"] = EditorTaskType.Missing.value
+            project_values["mediaType"] = MediaType.Audio.value
+            projects.add(Project(self.client, project_values))
+        return projects

libs/labelbox/tests/unit/schema/test_user_group.py

@@ -1,9 +1,11 @@
 import pytest
+from collections import defaultdict
 from unittest.mock import MagicMock
 from labelbox import Client
 from labelbox.exceptions import ResourceCreationError
+from labelbox.schema.project import Project
 from labelbox.schema.user import User
-from labelbox.schema.user_group import UserGroup, UserGroupColor, UserGroupUser, UserGroupProject
+from labelbox.schema.user_group import UserGroup, UserGroupColor
 
 
 class TestUserGroupColor:
@@ -20,40 +22,6 @@ def test_user_group_color_values(self):
         assert UserGroupColor.GRAY.value == "B8C4D3"
 
 
-class TestUserGroupUser:
-
-    def test_user_group_user_attributes(self):
-        user = UserGroupUser(id="user_id", email="test@example.com")
-        assert user.id == "user_id"
-        assert user.email == "test@example.com"
-
-    def test_user_group_user_equality(self):
-        user1 = UserGroupUser(id="user_id", email="test@example.com")
-        user2 = UserGroupUser(id="user_id", email="test@example.com")
-        assert user1 == user2
-
-    def test_user_group_user_hash(self):
-        user = UserGroupUser(id="user_id", email="test@example.com")
-        assert hash(user) == hash("user_id")
-
-
-class TestUserGroupProject:
-
-    def test_user_group_project_attributes(self):
-        project = UserGroupProject(id="project_id", name="Test Project")
-        assert project.id == "project_id"
-        assert project.name == "Test Project"
-
-    def test_user_group_project_equality(self):
-        project1 = UserGroupProject(id="project_id", name="Test Project")
-        project2 = UserGroupProject(id="project_id", name="Test Project")
-        assert project1 == project2
-
-    def test_user_group_project_hash(self):
-        project = UserGroupProject(id="project_id", name="Test Project")
-        assert hash(project) == hash("project_id")
-
-
 class TestUserGroup:
 
     def setup_method(self):
@@ -233,9 +201,13 @@ def test_create(self):
         group = self.group
         group.name = "New Group"
         group.color = UserGroupColor.PINK
-        group.users = {UserGroupUser(id="user_id", email="test@example.com")}
+        user_values = defaultdict(lambda: None)
+        user_values["id"] = "user_id"
+        user_values["email"] = "test@example.com"
+        group.users = {User(self.client, user_values)}
+        project_values = defaultdict(lambda: None)
         group.projects = {
-            UserGroupProject(id="project_id", name="Test Project")
+            Project(self.client, {id="project_id", name="Test Project", qu})
         }
 
         self.client.execute.return_value = {