add archive, batch.project, changelog

Author: farkob

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+# Version 3.20.0 (2022-04-27)
+## Added
+* Batches in a project can be retrieved with `project.batches()`
+* Added `Batch.remove_queued_data_rows()` to cancel remaining data rows in batch
+* Added `Batch.export_data_rows()` which returns `DataRow`s for a batch
+
+## Fix
+* `batch.project()` now works
+
 # Version 3.19.1 (2022-04-14)
 ## Fix
 * `create_data_rows` and `create_data_rows_sync` now uploads the file with a mimetype

labelbox/__init__.py

@@ -1,5 +1,5 @@
 name = "labelbox"
-__version__ = "3.19.1"
+__version__ = "3.20.0"
 
 import sys
 import warnings

labelbox/schema/batch.py

@@ -1,7 +1,8 @@
 from typing import Generator
 from labelbox.orm.db_object import DbObject, experimental
+from labelbox.orm import query
 from labelbox.orm.model import Entity, Field, Relationship
-from labelbox.exceptions import LabelboxError
+from labelbox.exceptions import LabelboxError, ResourceNotFoundError
 from io import StringIO
 import ndjson
 import requests
@@ -30,9 +31,49 @@ class Batch(DbObject):
     size = Field.Int("size")
 
     # Relationships
-    project = Relationship.ToOne("Project")
     created_by = Relationship.ToOne("User")
 
+    def __init__(self, client, project_id, *args, **kwargs):
+        super().__init__(client, *args, **kwargs)
+        self.project_id = project_id
+
+    def project(self) -> Entity.Project:
+        """ Returns Project which this Batch belongs to
+
+        Raises:
+            LabelboxError: if the project is not found
+        """
+        query_str = """query getProjectPyApi($projectId: ID!) {
+            project(
+                where: {id: $projectId}){
+                    %s
+                }}""" % query.results_query_part(Entity.Project)
+        params = {"projectId": self.project_id}
+        response = self.client.execute(query_str, params)
+
+        if response is None:
+            raise ResourceNotFoundError(Entity.Project, params)
+
+        return Entity.Project(self.client, response["project"])
+
+    def remove_queued_data_rows(self) -> None:
+        """ Removes remaining queued data rows from the batch and labeling queue.
+
+        Args:
+            batch (Batch): Batch to remove queued data rows from
+        """
+
+        project_id_param = "projectId"
+        batch_id_param = "batchId"
+        self.client.execute("""mutation ArchiveBatchPyApi($%s: ID!, $%s: ID!) {
+            project(where: {id: $%s}) { archiveBatch(batchId: $%s) { id archivedAt } }
+        }""" % (project_id_param, batch_id_param, project_id_param,
+                batch_id_param), {
+                    project_id_param: self.project_id,
+                    batch_id_param: self.uid
+                },
+                            experimental=True)
+
     def export_data_rows(self, timeout_seconds=120) -> Generator:
         """ Returns a generator that produces all data rows that are currently
         in this batch.

labelbox/schema/project.py

@@ -609,7 +609,7 @@ def create_batch(self, name: str, data_rows: List[str], priority: int = 5):
                                   experimental=True)["project"][method]
 
         res['size'] = len(dr_ids)
-        return Entity.Batch(self.client, res)
+        return Entity.Batch(self.client, self.uid, res)
 
     def _update_queue_mode(self,
                            mode: "Project.QueueMode") -> "Project.QueueMode":
@@ -851,17 +851,17 @@ def batches(self) -> PaginatedCollection:
             project(where: {id: $%s}) {id
             batches(after: $from, first: $first) { nodes { %s } pageInfo { endCursor }}}}
         """ % (id_param, id_param, query.results_query_part(Entity.Batch))
-        return PaginatedCollection(self.client,
-                                   query_str, {id_param: self.uid},
-                                   ['project', 'batches', 'nodes'],
-                                   Entity.Batch,
-                                   cursor_path={
-                                       'project': None,
-                                       'batches': None,
-                                       'pageInfo': None,
-                                       'endCursor': None
-                                   },
-                                   experimental=True)
+        return PaginatedCollection(
+            self.client,
+            query_str, {id_param: self.uid}, ['project', 'batches', 'nodes'],
+            lambda client, res: Entity.Batch(client, self.uid, res),
+            cursor_path={
+                'project': None,
+                'batches': None,
+                'pageInfo': None,
+                'endCursor': None
+            },
+            experimental=True)
 
     def upload_annotations(
             self,

tests/integration/test_batch.py

@@ -18,6 +18,19 @@ def big_dataset(dataset: Dataset):
     yield dataset
 
 
+@pytest.fixture
+def small_dataset(dataset: Dataset):
+    task = dataset.create_data_rows([
+        {
+            "row_data": IMAGE_URL,
+            "external_id": "my-image"
+        },
+    ] * 3)
+    task.wait_till_done()
+
+    yield dataset
+
+
 def test_create_batch(configured_project: Project, big_dataset: Dataset):
     configured_project.update(queue_mode=Project.QueueMode.Batch)
 
@@ -37,11 +50,30 @@ def test_export_data_rows(configured_project: Project, dataset: Dataset):
     ] * n_data_rows)
     task.wait_till_done()
 
-    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+    data_rows = [dr.uid for dr in list(dataset.export_data_rows())]
     batch = configured_project.create_batch("batch test", data_rows)
 
     result = list(batch.export_data_rows())
     exported_data_rows = [dr.uid for dr in result]
 
     assert len(result) == n_data_rows
     assert set(data_rows) == set(exported_data_rows)
+
+
+def test_archive_batch(configured_project: Project, small_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    batch = configured_project.create_batch("batch to archive", data_rows)
+    batch.remove_queued_data_rows()
+    exported_data_rows = list(batch.export_data_rows())
+
+    assert len(exported_data_rows) == 0
+
+
+def test_batch_project(configured_project: Project, small_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    batch = configured_project.create_batch(
+        "batch to test project relationship", data_rows)
+    project_from_batch = batch.project()
+
+    assert project_from_batch.uid == configured_project.uid
+    assert project_from_batch.name == configured_project.name