WIP batches

Author: gdj0nes

labelbox/client.py

@@ -1,27 +1,25 @@
-# type: ignore
-from datetime import datetime, timezone
 import json
-from typing import List, Dict
-from collections import defaultdict
-
 import logging
 import mimetypes
 import os
+from collections import defaultdict
+from datetime import datetime, timezone
+from typing import List, Dict
 
-from google.api_core import retry
 import requests
 import requests.exceptions
+from google.api_core import retry
 
 import labelbox.exceptions
-from labelbox import utils
 from labelbox import __version__ as SDK_VERSION
+from labelbox import utils
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Entity
 from labelbox.pagination import PaginatedCollection
+from labelbox.schema import role
 from labelbox.schema.data_row_metadata import DataRowMetadataOntology
 from labelbox.schema.iam_integration import IAMIntegration
-from labelbox.schema import role
 from labelbox.schema.ontology import Tool, Classification
 
 logger = logging.getLogger(__name__)
@@ -354,7 +352,7 @@ def upload_data(self,
             data=request_data,
             files={
                 "1": (filename, content, content_type) if
-                     (filename and content_type) else content
+                (filename and content_type) else content
             })
 
         if response.status_code == 502:
@@ -518,7 +516,7 @@ def _create(self, db_object_type, data):
         # Also convert Labelbox object values to their UIDs.
         data = {
             db_object_type.attribute(attr) if isinstance(attr, str) else attr:
-            value.uid if isinstance(value, DbObject) else value
+                value.uid if isinstance(value, DbObject) else value
             for attr, value in data.items()
         }
 
@@ -702,8 +700,8 @@ def get_data_row_ids_for_external_ids(
         for i in range(0, len(external_ids), max_ids_per_request):
             for row in self.execute(
                     query_str,
-                {'externalId_in': external_ids[i:i + max_ids_per_request]
-                })['externalIdsToDataRowIds']:
+                    {'externalId_in': external_ids[i:i + max_ids_per_request]
+                     })['externalIdsToDataRowIds']:
                 result[row['externalId']].append(row['dataRowId'])
         return result
 
@@ -896,3 +894,57 @@ def create_feature_schema(self, normalized):
         # But the features are the same so we just grab the feature schema id
         res['id'] = res['normalized']['featureSchemaId']
         return Entity.FeatureSchema(self, res)
+
+    def get_batch(self, batch_id: str):
+        """Gets a single Batch using its ID
+
+        Args:
+            batch_id: Id of the batch
+
+        Returns:
+            The sought Batch
+        """
+
+        return self._get_single(Entity.Batch, batch_id)
+
+    def create_batch(self, name: str, project, data_rows: List[str], priority: int):
+        """Create a batch of data rows to send to a project
+
+        >>> data_rows = ['<data-row-id>', ...]
+        >>> project = client.get("<project-id>")
+        >>> client.create_batch(name="low-confidence-images", project=project, data_rows=data_rows)
+
+        Args:
+            name: Descriptive name for the batch, must be unique per project
+            project: The project to send the batch to
+            data_rows: A list of data rows ids
+            priority: the default priority for the datarows, lowest priority by default
+
+        Returns:
+            The created batch
+        """
+
+        if isinstance(project, Entity.Project):
+            project_id = project.uid
+        elif isinstance(project, str):
+            project_id = project
+        else:
+            raise ValueError("You must pass a project id or a Project")
+
+        data_row_ids = []
+        for dr in data_rows:
+            pass
+
+        query_str = """mutation createBatchPyApi($name: String!, $dataRowIds: [ID!]!, $priority: Int!){
+            createBatch(){
+                    %s
+                }
+            }"""
+
+        result = self.execute(query_str, {
+            "name": name,
+            "projectId": project_id,
+            "dataRowIds": data_row_ids,
+            "priority": priority
+        })
+        return Entity.Batch(self, result['createModel'])

labelbox/schema/batch.py

@@ -0,0 +1,39 @@
+from typing import Dict, List
+
+from labelbox.orm.db_object import DbObject
+from labelbox.orm.model import Field, Relationship
+
+
+class Batch(DbObject):
+    """ A Batch is a group of data rows submitted to a project for labeling
+
+    Attributes:
+        name (str)
+        created_at (datetime)
+        updated_at (datetime)
+        deleted (bool)
+
+        project (Relationship): `ToOne` relationship to Project
+        created_by (Relationship): `ToOne` relationship to User
+
+    """
+    name = Field.String("name")
+    created_at = Field.DateTime("created_at")
+    updated_at = Field.DateTime("updated_at")
+    deleted = Field.Boolean()
+    size = Field.Int("size")
+
+    # Relationships
+    project = Relationship.ToOne("Project")
+    created_by = Relationship.ToOne("User")
+
+    def export_data_rows(self) -> List[Dict]:
+        """Get the data rows associated with a batch"""
+
+        gql = """query 
+        
+        """
+
+        self.client.execute()
+
+        return

labelbox/schema/project.py

@@ -2,11 +2,10 @@
 import json
 import logging
 import time
-import warnings
 from collections import namedtuple
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Dict, Union, Iterable, List, Optional
+from typing import Dict, Union, Iterable, Optional
 from urllib.parse import urlparse
 
 import ndjson
@@ -35,8 +34,6 @@
 
 logger = logging.getLogger(__name__)
 
-MAX_QUEUE_BATCH_SIZE = 1000
-
 
 class QueueMode(enum.Enum):
     Batch = "Batch"
@@ -154,6 +151,10 @@ def labels(self, datasets=None, order_by=None):
         return PaginatedCollection(self.client, query_str, {id_param: self.uid},
                                    ["project", "labels"], Label)
 
+    def batches(self):
+        """Returns a generator of batches that are queued for the project"""
+        # TODO
+
     def export_queued_data_rows(self, timeout_seconds=120):
         """ Returns all data rows that are currently enqueued for this project.
 
@@ -484,55 +485,6 @@ def setup(self, labeling_frontend, labeling_frontend_options):
         timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
         self.update(setup_complete=timestamp)
 
-    def queue(self, data_row_ids: List[str]):
-        """Add Data Rows to the Project queue"""
-
-        method = "submitBatchOfDataRows"
-        return self._post_batch(method, data_row_ids)
-
-    def dequeue(self, data_row_ids: List[str]):
-        """Remove Data Rows from the Project queue"""
-
-        method = "removeBatchOfDataRows"
-        return self._post_batch(method, data_row_ids)
-
-    def _post_batch(self, method, data_row_ids: List[str]):
-        """Post batch methods"""
-
-        if self.queue_mode() != QueueMode.Batch:
-            raise ValueError("Project must be in batch mode")
-
-        if len(data_row_ids) > MAX_QUEUE_BATCH_SIZE:
-            raise ValueError(
-                f"Batch exceeds max size of {MAX_QUEUE_BATCH_SIZE}, consider breaking it into parts"
-            )
-
-        query = """mutation %sPyApi($projectId: ID!, $dataRowIds: [ID!]!) {
-              project(where: {id: $projectId}) {
-                %s(data: {dataRowIds: $dataRowIds}) {
-                  dataRows {
-                    dataRowId
-                    error
-                  }
-                }
-              }
-            }
-        """ % (method, method)
-
-        res = self.client.execute(query, {
-            "projectId": self.uid,
-            "dataRowIds": data_row_ids
-        })["project"][method]["dataRows"]
-
-        # TODO: figure out error messaging
-        if len(data_row_ids) == len(res):
-            raise ValueError("No dataRows were submitted successfully")
-
-        if len(data_row_ids) > 0:
-            warnings.warn("Some Data Rows were not submitted successfully")
-
-        return res
-
     def _update_queue_mode(self, mode: QueueMode) -> QueueMode:
 
         if self.queue_mode() == mode:
@@ -846,7 +798,7 @@ class LabelingParameterOverride(DbObject):
 
 LabelerPerformance = namedtuple(
     "LabelerPerformance", "user count seconds_per_label, total_time_labeling "
-    "consensus average_benchmark_agreement last_activity_time")
+                          "consensus average_benchmark_agreement last_activity_time")
 LabelerPerformance.__doc__ = (
     "Named tuple containing info about a labeler's performance.")