Support Project.labels filting on Dataset

Author: florijanstamenkovic

labelbox/orm/model.py

@@ -194,11 +194,7 @@ def __init__(self, relationship_type, destination_type_name,
 
     @property
     def destination_type(self):
-        for t in Entity.subclasses():
-            if t.type_name() == self.destination_type_name:
-                return t
-        raise LabelboxError("Failed to find Entity for name: %s" %
-                            self.destination_type_name)
+        return Entity.named(self.destination_type_name)
 
 
 class Entity:
@@ -279,3 +275,20 @@ def subclasses(cls):
             yield subclass
             for subsub in subclass.subclasses():
                 yield subsub
+
+    @classmethod
+    def named(cls, name):
+        """ Returns an Entity (direct or indirect subclass of `Entity`) that
+        whose class name is equal to `name`.
+
+        Args:
+            name (str): Name of the sought entity, for example "Project".
+        Return:
+            An Entity subtype that has the given `name`.
+        Raises:
+            LabelboxError: if there is no such class.
+        """
+        for t in Entity.subclasses():
+            if t.type_name() == name:
+                return t
+        raise LabelboxError("Failed to find Entity for name: %s" % name)

labelbox/orm/query.py

@@ -374,7 +374,7 @@ def unset_labeling_parameter_overrides(project, data_rows):
     return query_str, {}
 
 
-def create_metadata(asset_type, meta_type, meta_value, data_row_id):
+def create_metadata(meta_type, meta_value, data_row_id):
     meta_type_param = "meta_type"
     meta_value_param = "meta_value"
     data_row_id_param = "data_row_id"
@@ -384,7 +384,8 @@ def create_metadata(asset_type, meta_type, meta_value, data_row_id):
             metaType: $%s metaValue: $%s dataRowId: $%s}) {%s}} """ % (
         meta_type_param, meta_value_param, data_row_id_param,
         meta_type_param, meta_value_param, data_row_id_param,
-        " ".join(field.graphql_name for field in asset_type.fields()))
+        " ".join(field.graphql_name for field
+                 in Entity.named("AssetMetadata").fields()))
     return query_str, {meta_type_param: meta_type,
                        meta_value_param: meta_value,
                        data_row_id_param: data_row_id}
@@ -484,6 +485,39 @@ def delete(db_object):
     return query_str, {id_param: db_object.uid}
 
 
+def project_labels(project, datasets, order_by):
+    """ Returns the query and params for getting a Project's labels
+    relationship. A non-standard relationship query is used to support
+    filtering on Datasets.
+    Args:
+        datasets (list or None): The datasets filter. If None it's
+            ignored.
+    Return:
+        (query_string, params)
+    """
+    label_entity = Entity.named("Label")
+
+    if datasets is not None:
+        where = " where:{dataRow: {dataset: {id_in: [%s]}}}" % ", ".join(
+            '"%s"' % dataset.uid for dataset in datasets)
+    else:
+        where = ""
+
+    if order_by is not None:
+        check_order_by_clause(label_entity, order_by)
+        order_by_str = "orderBy: %s_%s" % (
+            order_by[0].graphql_name, order_by[1].name.upper())
+    else:
+        order_by_str = ""
+
+    query_str = """query GetProjectLabelsPyApi($project_id: ID!)
+        {project (where: {id: $project_id})
+            {labels (skip: %%d first: %%d%s%s) {%s}}}""" % (
+        where, order_by_str, " ".join(f.graphql_name
+                                      for f in label_entity.fields()))
+    return query_str, {"project_id": project.uid}
+
+
 def export_labels():
     """ Returns the query and ID param for exporting a Project's
     labels.
@@ -521,9 +555,10 @@ def bulk_delete(db_objects, use_where_clause):
     return query_str, {}
 
 
-def create_webhook(entity, topics, url, secret, project):
+def create_webhook(topics, url, secret, project):
     project_str = "" if project is None else ("project:{id:\"%s\"}," % project.uid)
-    fields_str = " ".join(field.graphql_name for field in entity.fields())
+    fields_str = " ".join(field.graphql_name for field
+                          in Entity.named("Webhook").fields())
 
     query_str = """mutation CreateWebhookPyApi {
         createWebhook(data:{%s topics:{set:[%s]}, url:"%s", secret:"%s" }){%s}

labelbox/schema.py

@@ -10,6 +10,7 @@
 from labelbox.orm.db_object import (DbObject, Updateable, Deletable,
                                     BulkDeletable)
 from labelbox.orm.model import Field, Relationship
+from labelbox.pagination import PaginatedCollection
 
 
 """ Client-side object type definitions. """
@@ -19,12 +20,6 @@
 
 
 class Project(DbObject, Updateable, Deletable):
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.labels.supports_filtering = False
-
     name = Field.String("name")
     description = Field.String("description")
     updated_at = Field.DateTime("updated_at")
@@ -41,7 +36,6 @@ def __init__(self, *args, **kwargs):
     labeling_frontend = Relationship.ToOne("LabelingFrontend")
     labeling_frontend_options = Relationship.ToMany(
         "LabelingFrontendOptions", False, "labeling_frontend_options")
-    labels = Relationship.ToMany("Label", True)
     labeling_parameter_overrides = Relationship.ToMany(
         "LabelingParameterOverride", False, "labeling_parameter_overrides")
     webhooks = Relationship.ToMany("Webhook", False)
@@ -72,6 +66,11 @@ def create_label(self, **kwargs):
         res = res["data"]["createLabel"]
         return Label(self.client, res)
 
+    def labels(self, datasets=None, order_by=None):
+        query_string, params = query.project_labels(self, datasets, order_by)
+        return PaginatedCollection(self.client, query_string, params,
+                                   ["project", "labels"], Label)
+
     def export_labels(self, timeout_seconds=60):
         """ Calls the server-side Label exporting that generates a JSON
         payload, and returns the URL to that payload.
@@ -311,8 +310,7 @@ def create_metadata(self, meta_type, meta_value):
         Return:
             AssetMetadata DB object.
         """
-        query_str, params = query.create_metadata(
-            AssetMetadata, meta_type, meta_value, self.uid)
+        query_str, params = query.create_metadata(meta_type, meta_value, self.uid)
         res = self.client.execute(query_str, params)
         return AssetMetadata(self.client, res["data"]["createAssetMetadata"])
 
@@ -454,8 +452,7 @@ class Webhook(DbObject):
 
     @staticmethod
     def create(client, topics, url, secret, project):
-        query_str, params = query.create_webhook(Webhook, topics, url, secret,
-                                                 project)
+        query_str, params = query.create_webhook(topics, url, secret, project)
         res = client.execute(query_str, params)
         return Webhook(client, res["data"]["createWebhook"])
 

tests/integration/test_label.py

@@ -80,27 +80,33 @@ def test_label_update(client, rand_gen):
 
 def test_label_filter_order(client, rand_gen):
     project = client.create_project(name=rand_gen(str))
-    dataset = client.create_dataset(name=rand_gen(str), projects=project)
-    data_row = dataset.create_data_row(row_data=IMG_URL)
+    dataset_1 = client.create_dataset(name=rand_gen(str), projects=project)
+    dataset_2 = client.create_dataset(name=rand_gen(str), projects=project)
+    data_row_1 = dataset_1.create_data_row(row_data=IMG_URL)
+    data_row_2 = dataset_2.create_data_row(row_data=IMG_URL)
 
-    l1 = project.create_label(data_row=data_row, label="l1", seconds_to_label=0.3)
-    l2 = project.create_label(data_row=data_row, label="l2", seconds_to_label=0.1)
+    l1 = project.create_label(data_row=data_row_1, label="l1",
+                              seconds_to_label=0.3)
+    l2 = project.create_label(data_row=data_row_2, label="l2",
+                              seconds_to_label=0.1)
 
     # Labels are not visible in the project immediately.
     time.sleep(10)
 
-    # Filtering is not supported
-    with pytest.raises(InvalidQueryError) as exc_info:
-        project.labels(where=Label.label=="l1")
-    assert exc_info.value.message == \
-        "Relationship Project.labels doesn't support filtering"
+    # Filtering supported on dataset
+    assert set(project.labels()) == {l1, l2}
+    assert set(project.labels(datasets=[])) == set()
+    assert set(project.labels(datasets=[dataset_1])) == {l1}
+    assert set(project.labels(datasets=[dataset_2])) == {l2}
+    assert set(project.labels(datasets=[dataset_1, dataset_2])) == {l1, l2}
 
     assert list(project.labels(order_by=Label.label.asc)) == [l1, l2]
     assert list(project.labels(order_by=Label.label.desc)) == [l2, l1]
     assert list(project.labels(order_by=Label.seconds_to_label.asc)) == [l2, l1]
     assert list(project.labels(order_by=Label.seconds_to_label.desc)) == [l1, l2]
 
-    dataset.delete()
+    dataset_1.delete()
+    dataset_2.delete()
     project.delete()
 
 

tools/db_object_doc_gen.py

@@ -11,11 +11,17 @@
 sys.path.append(sys.argv[1])
 
 
-from labelbox import (Project, Dataset, DataRow, Label, User, Organization, Task,
-                      LabelingFrontend, Webhook)
+from labelbox import (Project, Dataset, DataRow, Label, User, Organization,
+                      Task, LabelingFrontend, Webhook)
 from labelbox.schema import Field, Relationship
 
 
+# Map {cls: [fields_info, ...]}
+more_fields = {}
+# Map {cls: [relationship_info, ...]}
+more_rels = {Project: [("labels", Label, "Many")]}
+
+
 for cls in (Project, Dataset, DataRow, Label, User, Organization, Task,
             LabelingFrontend, Webhook):
     print("")
@@ -34,3 +40,6 @@
                          if isinstance(r, Relationship)):
         print(relationship.name, "|", relationship.destination_type_name, "|",
               relationship.relationship_type.name[2:])
+
+    for name, dest_type, card in more_rels.get(cls, []):
+        print(name, "|", dest_type.type_name(), "|", card)