Add Dataset.data_row_for_external_id

Author: florijanstamenkovic

labelbox/schema.py

@@ -5,7 +5,7 @@
 import os
 import time
 
-from labelbox.exceptions import InvalidQueryError
+from labelbox.exceptions import InvalidQueryError, ResourceNotFoundError
 from labelbox.orm import query
 from labelbox.orm.db_object import (DbObject, Updateable, Deletable,
                                     BulkDeletable)
@@ -205,9 +205,11 @@ def create_data_rows(self, items):
 
         Args:
             items (iterable of (dict or str)): See above for details.
+
         Return:
             Task representing the data import on the server side. The Task
             can be used for inspecting task progress and waiting until it's done.
+
         Raise:
             InvalidQueryError: if the `items` parameter does not conform to
                 the specification above.
@@ -278,6 +280,32 @@ def convert_item(item):
         task._user = user
         return task
 
+    def data_row_for_external_id(self, external_id):
+        """ Convenience method for getting a single `DataRow` belonging to this
+        `Dataset` that has the given `external_id`.
+
+        Args:
+            external_id (str): External ID of the sought `DataRow`.
+
+        Returns:
+            A single `DataRow` with the given ID.
+
+        Raises:
+            labelbox.exceptions.ResourceNotFoundError: if there is no `DataRow`
+                in this `DataSet` with the given external ID, or if there are
+                multiple `DataRows` for it.
+        """
+        where = DataRow.external_id==external_id
+
+        data_rows = self.data_rows(where=where)
+        # Get at most two data_rows.
+        data_rows = [row for row, _ in zip(data_rows, range(2))]
+
+        if len(data_rows) != 1:
+            raise ResourceNotFoundError(DataRow, where)
+
+        return data_rows[0]
+
 
 class DataRow(DbObject, Updateable, BulkDeletable):
     external_id = Field.String("external_id")

tests/integration/test_dataset.py

@@ -4,6 +4,9 @@
 from labelbox.exceptions import ResourceNotFoundError
 
 
+IMG_URL = "https://picsum.photos/200/300"
+
+
 def test_dataset(client, rand_gen):
     before = list(client.get_datasets())
     for o in before:
@@ -55,3 +58,22 @@ def test_dataset_filtering(client):
 
     d1.delete()
     d2.delete()
+
+
+def test_get_data_row_for_external_id(client, rand_gen):
+    dataset = client.create_dataset(name=rand_gen(str))
+    external_id = rand_gen(str)
+
+    with pytest.raises(ResourceNotFoundError):
+        data_row = dataset.data_row_for_external_id(external_id)
+
+    data_row = dataset.create_data_row(row_data=IMG_URL, external_id=external_id)
+
+    found = dataset.data_row_for_external_id(external_id)
+    assert found.uid == data_row.uid
+    assert found.external_id == external_id
+
+    second = dataset.create_data_row(row_data=IMG_URL, external_id=external_id)
+
+    with pytest.raises(ResourceNotFoundError):
+        data_row = dataset.data_row_for_external_id(external_id)