Use Entity metaclass for schema type registration.

Author: florijanstamenkovic

labelbox/orm/model.py

@@ -194,7 +194,7 @@ def __init__(self, relationship_type, destination_type_name,
 
     @property
     def destination_type(self):
-        return Entity.named(self.destination_type_name)
+        return getattr(Entity, self.destination_type_name)
 
     def __str__(self):
         return self.name
@@ -203,7 +203,18 @@ def __repr__(self):
         return "<Relationship: %r>" % self.name
 
 
-class Entity:
+class EntityMeta(type):
+    """ Entity metaclass. Registers Entity subclasses as attributes
+    of the Entity class object so they can be referenced for example like:
+        Entity.Project.
+    """
+    def __init__(cls, clsname, superclasses, attributedict):
+        super().__init__(clsname, superclasses, attributedict)
+        if clsname != "Entity":
+            setattr(Entity, clsname, cls)
+
+
+class Entity(metaclass=EntityMeta):
     """ An entity that contains fields and relationships. Base class
     for DbObject (which is base class for concrete schema classes). """
 
@@ -279,27 +290,3 @@ def type_name(cls):
             Project, DataRow, ...
         """
         return cls.__name__.split(".")[-1]
-
-
-    @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.
-        """
-        def subclasses(cls):
-            for subclass in cls.__subclasses__():
-                yield subclass
-                for subsub in subclasses(subclass):
-                    yield subsub
-
-        for t in subclasses(Entity):
-            if t.type_name() == name:
-                return t
-        raise LabelboxError("Failed to find Entity for name: %s" % name)

labelbox/schema/data_row.py

@@ -51,10 +51,10 @@ def create_metadata(self, meta_type, meta_value):
                 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,
-            query.results_query_part(Entity.named("AssetMetadata")))
+            query.results_query_part(Entity.AssetMetadata))
 
         res = self.client.execute(
             query_str, {meta_type_param: meta_type, meta_value_param: meta_value,
                         data_row_id_param: self.uid})
-        return Entity.named("AssetMetadata")(
+        return Entity.AssetMetadata(
             self.client, res["data"]["createAssetMetadata"])

labelbox/schema/dataset.py

@@ -38,7 +38,7 @@ def create_data_row(self, **kwargs):
                 any of the field names given in `data`.
 
         """
-        DataRow = Entity.named("DataRow")
+        DataRow = Entity.DataRow
         if DataRow.row_data.name not in kwargs:
             raise InvalidQueryError(
                 "DataRow.row_data missing when creating DataRow.")
@@ -80,8 +80,7 @@ def create_data_rows(self, items):
                 a DataRow.
         """
         file_upload_thread_count = 20
-        DataRow = Entity.named("DataRow")
-        Task = Entity.named("Task")
+        DataRow = Entity.DataRow
 
         def upload_if_necessary(item):
             if isinstance(item, str):
@@ -136,12 +135,12 @@ def convert_item(item):
         # Fetch and return the task.
         task_id = res["taskId"]
         user = self.client.get_user()
-        task = list(user.created_tasks(where=Task.uid == task_id))
+        task = list(user.created_tasks(where=Entity.Task.uid == task_id))
         # Cache user in a private variable as the relationship can't be
         # resolved due to server-side limitations (see Task.created_by)
         # for more info.
         if len(task) != 1:
-            raise ResourceNotFoundError(Task, task_id)
+            raise ResourceNotFoundError(Entity.Task, task_id)
         task = task[0]
         task._user = user
         return task
@@ -161,7 +160,7 @@ def data_row_for_external_id(self, external_id):
                 in this `DataSet` with the given external ID, or if there are
                 multiple `DataRows` for it.
         """
-        DataRow = Entity.named("DataRow")
+        DataRow = Entity.DataRow
         where = DataRow.external_id==external_id
 
         data_rows = self.data_rows(where=where)

labelbox/schema/label.py

@@ -38,10 +38,9 @@ def create_review(self, **kwargs):
             Review attributes. At a minimum a `Review.score` field
             value must be provided.
         """
-        Review = Entity.named("Review")
-        kwargs[Review.label.name] = self
-        kwargs[Review.project.name] = self.project()
-        return self.client._create(Review, kwargs)
+        kwargs[Entity.Review.label.name] = self
+        kwargs[Entity.Review.project.name] = self.project()
+        return self.client._create(Entity.Review, kwargs)
 
     def create_benchmark(self):
         """ Creates a Benchmark for this Label.
@@ -52,7 +51,7 @@ def create_benchmark(self):
         query_str = """mutation CreateBenchmarkPyApi($%s: ID!) {
             createBenchmark(data: {labelId: $%s}) {%s}} """ % (
                 label_id_param, label_id_param,
-                query.results_query_part(Entity.named("Benchmark")))
+                query.results_query_part(Entity.Benchmark))
         res = self.client.execute(query_str, {label_id_param: self.uid})
         res = res["data"]["createBenchmark"]
-        return Entity.named("Benchmark")(self.client, res)
+        return Entity.Benchmark(self.client, res)

labelbox/schema/project.py

@@ -54,7 +54,7 @@ def create_label(self, **kwargs):
         # about them. At the same time they're connected to a Label at
         # label creation in a non-standard way (connect via name).
 
-        Label = Entity.named("Label")
+        Label = Entity.Label
         kwargs[Label.project] = self
         data = {Label.attribute(attr) if isinstance(attr, str) else attr:
                 value.uid if isinstance(value, DbObject) else value
@@ -69,7 +69,7 @@ def create_label(self, **kwargs):
         return Label(self.client, res)
 
     def labels(self, datasets=None, order_by=None):
-        Label = Entity.named("Label")
+        Label = Entity.Label
 
         if datasets is not None:
             where = " where:{dataRow: {dataset: {id_in: [%s]}}}" % ", ".join(
@@ -135,10 +135,10 @@ def labeler_performance(self):
                     count user {%s} secondsPerLabel totalTimeLabeling consensus
                     averageBenchmarkAgreement lastActivityTime}
             }}""" % (project_id_param, project_id_param,
-                     query.results_query_part(Entity.named("User")))
+                     query.results_query_part(Entity.User))
 
         def create_labeler_performance(client, result):
-            result["user"] = Entity.named("User")(client, result["user"])
+            result["user"] = Entity.User(client, result["user"])
             result["lastActivityTime"] = datetime.fromtimestamp(
                 result["lastActivityTime"] / 1000, timezone.utc)
             return LabelerPerformance(**{utils.snake_case(key): value
@@ -155,7 +155,7 @@ def review_metrics(self, net_score):
         Return:
             int, aggregation count of reviews for given net_score.
         """
-        if net_score not in (None,) + tuple(Entity.named("Review").NetScore):
+        if net_score not in (None,) + tuple(Entity.Review.NetScore):
             raise InvalidQueryError("Review metrics net score must be either None "
                                     "or one of Review.NetScore values")
         project_id_param = "project_id"
@@ -179,7 +179,7 @@ def setup(self, labeling_frontend, labeling_frontend_options):
         if not isinstance(labeling_frontend_options, str):
             labeling_frontend_options = json.dumps(labeling_frontend_options)
 
-        LFO = Entity.named("LabelingFrontendOptions")
+        LFO = Entity.LabelingFrontendOptions
         labeling_frontend_options = self.client._create(
             LFO, {LFO.project: self, LFO.labeling_frontend: labeling_frontend,
                   LFO.customization_options: labeling_frontend_options,

labelbox/schema/webhook.py

@@ -48,7 +48,7 @@ def create(client, topics, url, secret, project):
         query_str = """mutation CreateWebhookPyApi {
             createWebhook(data:{%s topics:{set:[%s]}, url:"%s", secret:"%s" }){%s}
         } """ % (project_str, " ".join(topics), url, secret,
-                query.results_query_part(Entity.named("Webhook")))
+                query.results_query_part(Entity.Webhook))
 
         res = client.execute(query_str)
         return Webhook(client, res["data"]["createWebhook"])
@@ -74,7 +74,7 @@ def update(self, topics=None, url=None, status=None):
         query_str = """mutation UpdateWebhookPyApi {
             updateWebhook(where: {id: "%s"} data:{%s}){%s}} """ % (
             self.uid, ", ".join(filter(None, (topics_str, url_str, status_str))),
-            query.results_query_part(Entity.named("Webhook")))
+            query.results_query_part(Entity.Webhook))
 
         res = self.client.execute(query_str)
         res = res["data"]["updateWebhook"]