diff --git a/labelbox/schema/batch.py b/labelbox/schema/batch.py index 9a30349aa..f14ea55d2 100644 --- a/labelbox/schema/batch.py +++ b/labelbox/schema/batch.py @@ -1,4 +1,5 @@ from typing import Generator, TYPE_CHECKING + from labelbox.orm.db_object import DbObject, experimental from labelbox.orm import query from labelbox.orm.model import Entity, Field, Relationship diff --git a/labelbox/schema/export_filters.py b/labelbox/schema/export_filters.py index 2fe6742fa..88af09f67 100644 --- a/labelbox/schema/export_filters.py +++ b/labelbox/schema/export_filters.py @@ -37,7 +37,11 @@ class SharedExportFilters(TypedDict): class ProjectExportFilters(SharedExportFilters): - pass + batch_id: Optional[str] + """ Batch id to export + Example: + >>> "clgo3lyax0000veeezdbu3ws4" + """ class DatasetExportFilters(SharedExportFilters): @@ -178,4 +182,12 @@ def _get_timezone() -> str: "operator": "is", "type": "data_row_id" }) + + batch_id = filters.get("batch_id") + if batch_id: + search_query.append({ + "ids": [batch_id], + "operator": "is", + "type": "batch" + }) return search_query diff --git a/labelbox/schema/project.py b/labelbox/schema/project.py index 153459bab..ff3f37e8f 100644 --- a/labelbox/schema/project.py +++ b/labelbox/schema/project.py @@ -442,6 +442,7 @@ def export_v2(self, "last_activity_at": None, "label_created_at": None, "data_row_ids": None, + "batch_id": None, }) mutation_name = "exportDataRowsInProject" diff --git a/tests/integration/test_batch.py b/tests/integration/test_batch.py index 62d4e906b..9df0a89e2 100644 --- a/tests/integration/test_batch.py +++ b/tests/integration/test_batch.py @@ -204,4 +204,4 @@ def test_delete_labels_with_templates(batch_project: Project, exported_data_rows = list(batch.export_data_rows()) res = batch.delete_labels(labels_as_template=True) exported_data_rows = list(batch.export_data_rows()) - assert len(exported_data_rows) == 5 + assert len(exported_data_rows) == 5 \ No newline at end of file diff --git a/tests/integration/test_project.py b/tests/integration/test_project.py index 253808fc7..6b0e889ce 100644 --- a/tests/integration/test_project.py +++ b/tests/integration/test_project.py @@ -1,5 +1,6 @@ import time import os +from typing import Tuple import uuid import pytest @@ -7,6 +8,8 @@ from labelbox import Project, LabelingFrontend, Dataset from labelbox.exceptions import InvalidQueryError +from labelbox.schema.data_row import DataRow +from labelbox.schema.label import Label from labelbox.schema.media_type import MediaType from labelbox.schema.queue_mode import QueueMode @@ -42,6 +45,43 @@ def test_project(client, rand_gen): assert project not in projects +def test_batch_project_export_v2( + configured_batch_project_with_label: Tuple[Project, Dataset, DataRow, + Label], + export_v2_test_helpers, dataset: Dataset, image_url: str): + project, dataset, *_ = configured_batch_project_with_label + + batch = list(project.batches())[0] + filters = { + "last_activity_at": ["2000-01-01 00:00:00", "2050-01-01 00:00:00"], + "label_created_at": ["2000-01-01 00:00:00", "2050-01-01 00:00:00"], + "batch_id": batch.uid, + } + params = { + "include_performance_details": True, + "include_labels": True, + "media_type_override": MediaType.Image + } + task_name = "test_batch_export_v2" + task = dataset.create_data_rows([ + { + "row_data": image_url, + "external_id": "my-image" + }, + ] * 2) + task.wait_till_done() + data_rows = [dr.uid for dr in list(dataset.export_data_rows())] + batch_one = f'batch one {uuid.uuid4()}' + + # This test creates two batches, only one batch should be exporter + # Creatin second batch that will not be used in the export due to the filter: batch_id + project.create_batch(batch_one, data_rows) + + task_results = export_v2_test_helpers.run_project_export_v2_task( + project, task_name=task_name, filters=filters, params=params) + assert (batch.size == len(task_results)) + + def test_project_export_v2(client, export_v2_test_helpers, configured_project_with_label, wait_for_data_row_processing):