Add logic to upload labels by project id to model run

Author: Val Brodsky

labelbox/schema/model_run.py

@@ -47,22 +47,41 @@ class Status(Enum):
         COMPLETE = "COMPLETE"
         FAILED = "FAILED"
 
-    def upsert_labels(self, label_ids, timeout_seconds=3600):
+    def upsert_labels(self,
+        label_ids: Optional[List[str]] = None,
+        project_id: Optional[str] = None, 
+        timeout_seconds=3600):
         """ Adds data rows and labels to a Model Run
         Args:
             label_ids (list): label ids to insert
+            project_id (string): project uuid, all project labels will be uploaded
             timeout_seconds (float): Max waiting time, in seconds.
         Returns:
             ID of newly generated async task
         """
 
-        if len(label_ids) < 1:
-            raise ValueError("Must provide at least one label id")
+        use_label_ids = label_ids is not None and len(label_ids) > 0
+        use_project_id = project_id is not None
 
+        if not use_label_ids and not use_project_id:
+            raise ValueError("Must provide at least one label id or a project id")
+
+        if use_label_ids and use_project_id:
+            raise ValueError("Must only one of label ids, project id")
+
+        if use_label_ids:
+            return self._upsert_labels_by_label_ids(label_ids)
+        else: # use_project_id
+            return self._upsert_labels_by_project_id(project_id)
+
+
+
+    def _upsert_labels_by_label_ids(self,
+        label_ids: List[str]):
         mutation_name = 'createMEAModelRunLabelRegistrationTask'
         create_task_query_str = """mutation createMEAModelRunLabelRegistrationTaskPyApi($modelRunId: ID!, $labelIds : [ID!]!) {
-          %s(where : { id : $modelRunId}, data : {labelIds: $labelIds})}
-          """ % (mutation_name)
+        %s(where : { id : $modelRunId}, data : {labelIds: $labelIds})}
+        """ % (mutation_name)
 
         res = self.client.execute(create_task_query_str, {
             'modelRunId': self.uid,
@@ -78,7 +97,31 @@ def upsert_labels(self, label_ids, timeout_seconds=3600):
             status_query_str, {'where': {
                 'id': task_id
             }})['MEALabelRegistrationTaskStatus'],
-                                     timeout_seconds=timeout_seconds)
+                                        timeout_seconds=timeout_seconds)
+
+    def _upsert_labels_by_project_id(self,
+        project_id: str):
+        mutation_name = 'createMEAModelRunProjectLabelRegistrationTask'
+        create_task_query_str = """mutation createMEAModelRunProjectLabelRegistrationTaskPyApi($modelRunId: ID!, $projectId : ID!) {
+        %s(where : { modelRunId : $modelRunId, projectId: $projectId}}
+        """ % (mutation_name)
+
+        res = self.client.execute(create_task_query_str, {
+            'modelRunId': self.uid,
+            'projectId': project_id
+        })
+        task_id = res[mutation_name]
+
+        status_query_str = """query MEALabelRegistrationTaskStatusPyApi($where: WhereUniqueIdInput!){
+            MEALabelRegistrationTaskStatus(where: $where) {status errorMessage}
+        }
+        """
+        return self._wait_until_done(lambda: self.client.execute(
+            status_query_str, {'where': {
+                'id': task_id
+            }})['MEALabelRegistrationTaskStatus'],
+                                        timeout_seconds=timeout_seconds)
+
 
     def upsert_data_rows(self,
                          data_row_ids=None,