PLT-1016 - Add methods for External Workforce management

docs/labelbox/external-workforce.rst

@@ -0,0 +1,6 @@
+ExternalWorkforce
+===============================================================================================
+
+.. automodule:: labelbox.schema.external_workforce
+   :members:
+   :show-inheritance:

libs/labelbox/src/labelbox/__init__.py

@@ -40,3 +40,4 @@
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.schema.ontology_kind import OntologyKind
 from labelbox.schema.project_overview import ProjectOverview, ProjectOverviewDetailed
+from labelbox.schema.external_workforce import ExternalWorkforce

libs/labelbox/src/labelbox/schema/__init__.py

@@ -25,4 +25,5 @@
 import labelbox.schema.identifiable
 import labelbox.schema.catalog
 import labelbox.schema.ontology_kind
+import labelbox.schema.external_workforce
 import labelbox.schema.project_overview

libs/labelbox/src/labelbox/schema/external_workforce.py

@@ -0,0 +1,12 @@
+from labelbox.pydantic_compat import BaseModel
+
+class ExternalWorkforce(BaseModel):
+    """
+    Represents an external workforce used in the Labelbox system.
+
+    Attributes:
+        id (str): The unique identifier of the external workforce.
+        name (str): The name of the external workforce.
+    """
+    id: str
+    name: str

libs/labelbox/src/labelbox/schema/project.py

@@ -42,6 +42,7 @@
 from labelbox.schema.task_queue import TaskQueue
 from labelbox.schema.ontology_kind import (EditorTaskType, OntologyKind)
 from labelbox.schema.project_overview import ProjectOverview, ProjectOverviewDetailed
+from labelbox.schema.external_workforce import ExternalWorkforce
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
@@ -1840,6 +1841,115 @@ def clone(self) -> "Project":
         """
         result = self.client.execute(mutation, {"projectId": self.uid})
         return self.client.get_project(result["cloneProject"]["id"])
+
+    def add_external_workforce(self, workforce_id: Union[str, ExternalWorkforce]) -> List[ExternalWorkforce]:
+
+        """
+        Add an external workforce (organization) to a project.
+
+        Args:
+            workforce_id (Union[str, ExternalWorkforce]): Organization id of the external workforce.
+
+        Returns:
+            List[ExternalWorkforce]: The remaining external workforces or an empty list if no external workforces are left.
+
+        Raises:
+            LabelboxError: If the external workforce with the given ID cannot be found.
+
+        Note:
+            This method adds an external workforce (organization) to the current project. 
+            The `workforce_id` parameter can be either a string representing the organization ID or an instance of the `ExternalWorkforce` class.
+        """
+        workforce_id = workforce_id.uid if isinstance(workforce_id, ExternalWorkforce) else workforce_id
+
+        mutation = """
+            mutation ShareProjectWithExternalOrganizationPyApi($projectId: ID!, $organizationId: ID!) {
+                shareProjectWithExternalOrganization(
+                    data: { projectId: $projectId, organizationId: $organizationId }
+                ) {
+                    id
+                    sharedWithOrganizations {
+                        id
+                        name
+                    }
+                    }
+            }
+        """
+
+        result = self.client.execute(mutation, {"projectId": self.uid, "organizationId": workforce_id})
+
+        if not result:
+            raise LabelboxError(f"Can't find External Workforce {workforce_id}")
+
+        return [ExternalWorkforce(**workforce_dic) 
+                    for workforce_dic 
+                    in result["shareProjectWithExternalOrganization"]["sharedWithOrganizations"]]
+
+
+    def remove_external_workforce(self, workforce_id: Union[str, ExternalWorkforce]) -> List[ExternalWorkforce]:
+        """Remove an external workforce (organization) from a project.
+
+        Args:
+            workforce_id (Union[str, ExternalWorkforce]): Organization id of the external workforce
+            or an instance of the ExternalWorkforce class.
+
+        Returns:
+            List[ExternalWorkforce]: The remaining external workforces or an empty list if no external workforces are left.
+
+        Raises:
+            LabelboxError: If the external workforce cannot be found.
+        """
+        workforce_id = workforce_id.uid if isinstance(workforce_id, ExternalWorkforce) else workforce_id
+
+        mutation = """
+            mutation UnshareProjectWithExternalOrganizationPyApi($projectId: ID!, $organizationId: ID!) {
+                unshareProjectWithExternalOrganization(
+                    data: { projectId: $projectId, organizationId: $organizationId }
+                ) {
+                    id
+                    sharedWithOrganizations {
+                        id
+                        name
+                    }
+                }
+            }
+            """
+
+        result = self.client.execute(mutation, {"projectId": self.uid, "organizationId": workforce_id})
+
+        if not result:
+            raise LabelboxError(f"Can't find External Workforce {workforce_id}")
+
+        return [ExternalWorkforce(**workforce_dic) 
+                for workforce_dic 
+                in result["unshareProjectWithExternalOrganization"]["sharedWithOrganizations"]]
+
+
+    def get_external_workforces(self) -> List[ExternalWorkforce]:
+        """List the external workforces (organizations) attached to a project
+
+        Args:
+            project_id: Id of the project to check Project to check
+
+        Returns:
+            list of dictionaries with id and name for each external workforce (aka organization)
+            or empty list if no external workforces were found
+        """
+
+        query = """
+            query GetProjectExternalOganizationsPyApi($projectId: ID!) {
+                project(where: { id: $projectId }) {
+                    id
+                    sharedWithOrganizations {
+                        id
+                        name
+                    }
+                }
+            }
+            """
+
+        result = self.client.execute(query, {"projectId": self.uid})["project"]["sharedWithOrganizations"]
+        return [ExternalWorkforce(**workforce_dic) for workforce_dic in result]  
 
 
 class ProjectMember(DbObject):

libs/labelbox/tests/integration/test_external_workforce.py

@@ -0,0 +1,41 @@
+from labelbox import Project, ExternalWorkforce
+import os
+import sys
+
+def _get_workforce_id_for_test() -> str:
+    # Basic function to provide an organization id for the test
+    # org id from https://labelbox.atlassian.net/wiki/spaces/PLT/pages/2110816271/How+to+labelbox-python+SDK+CI+Tests
+    current_version = sys.version_info[:2] 
+
+    if os.environ['LABELBOX_TEST_ENVIRON'] == "staging":
+        return "cltp16p2v04dr07ywfgqxf23u" if current_version == (3, 8) else "clum46jsb00jp07z338dh1gvb"
+    else:
+        return "cltp1fral01dh07009zsng3zs" if current_version == (3, 8) else "ckcz6bubudyfi0855o1dt1g9s"
+
+
+def test_add_external_workforce(project: Project):
+    workforce_id = _get_workforce_id_for_test()
+
+    external_workforces = project.add_external_workforce(workforce_id)
+    assert len(external_workforces) == 1
+    assert isinstance(external_workforces[0], ExternalWorkforce)
+
+
+def test_get_external_workforces(project: Project):
+    workforce_id = _get_workforce_id_for_test()
+
+    external_workforces = project.add_external_workforce(workforce_id)
+
+    external_workforces = project.get_external_workforces()
+    assert len(external_workforces) == 1
+    assert isinstance(external_workforces[0], ExternalWorkforce)
+
+
+def test_remove_external_workforce(project: Project):    
+    workforce_id = _get_workforce_id_for_test()
+    external_workforces = project.add_external_workforce(workforce_id)
+
+    external_workforces = project.remove_external_workforce(workforce_id)
+    assert len(external_workforces) == 0
+
+