update to clean up some dependencies and utilize Entity more

Author: jtsodapop

labelbox/orm/model.py

@@ -343,6 +343,10 @@ class Entity(metaclass=EntityMeta):
         IAMIntegration: Type[labelbox.IAMIntegration]
         LabelingFrontendOptions: Type[labelbox.LabelingFrontendOptions]
         Label: Type[labelbox.Label]
+        MEAPredictionImport: Type[labelbox.MEAPredictionImport]
+        Invite: Type[labelbox.Invite]
+        InviteLimit: Type[labelbox.InviteLimit]
+        ProjectRole: Type[labelbox.ProjectRole]
 
     @classmethod
     def _attributes_of_type(cls, attr_type):

labelbox/schema/bulk_import_request.py

@@ -12,15 +12,19 @@
 from pydantic import BaseModel, validator
 from typing_extensions import Literal
 from typing import (Any, List, Optional, BinaryIO, Dict, Iterable, Tuple, Union,
-                    Type, Set)
+                    Type, Set, TYPE_CHECKING)
 
-import labelbox
+from labelbox import exceptions as lb_exceptions
+from labelbox.orm.model import Entity
 from labelbox import utils
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field, Relationship
 from labelbox.schema.enums import BulkImportRequestState
 
+if TYPE_CHECKING:
+    from labelbox import Project
+
 NDJSON_MIME_TYPE = "application/x-ndjson"
 logger = logging.getLogger(__name__)
 
@@ -67,7 +71,7 @@ def _send_create_file_command(
     response = client.execute(data=request_data, files={file_name: file_data})
 
     if not response.get("createBulkImportRequest", None):
-        raise labelbox.exceptions.LabelboxError(
+        raise lb_exceptions.LabelboxError(
             "Failed to create BulkImportRequest, message: %s" %
             response.get("errors", None) or response.get("error", None))
 
@@ -192,9 +196,8 @@ def wait_until_done(self, sleep_time_seconds: int = 5) -> None:
             self.__exponential_backoff_refresh()
 
     @backoff.on_exception(
-        backoff.expo,
-        (labelbox.exceptions.ApiLimitError, labelbox.exceptions.TimeoutError,
-         labelbox.exceptions.NetworkError),
+        backoff.expo, (lb_exceptions.ApiLimitError, lb_exceptions.TimeoutError,
+                       lb_exceptions.NetworkError),
         max_tries=10,
         jitter=None)
     def __exponential_backoff_refresh(self) -> None:
@@ -387,7 +390,7 @@ def delete(self) -> None:
 
 
 def _validate_ndjson(lines: Iterable[Dict[str, Any]],
-                     project: "labelbox.Project") -> None:
+                     project: "Project") -> None:
     """
     Client side validation of an ndjson object.
 
@@ -413,12 +416,12 @@ def _validate_ndjson(lines: Iterable[Dict[str, Any]],
             annotation.validate_instance(feature_schemas)
             uuid = str(annotation.uuid)
             if uuid in uids:
-                raise labelbox.exceptions.UuidError(
+                raise lb_exceptions.UuidError(
                     f'{uuid} already used in this import job, '
                     'must be unique for the project.')
             uids.add(uuid)
         except (pydantic.ValidationError, ValueError, TypeError, KeyError) as e:
-            raise labelbox.exceptions.MALValidationError(
+            raise lb_exceptions.MALValidationError(
                 f"Invalid NDJson on line {idx}") from e
 
 

labelbox/schema/data_row.py

@@ -1,11 +1,13 @@
 import logging
 from datetime import datetime
-from typing import List, Dict, Union
+from typing import List, Dict, Union, TYPE_CHECKING
 
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Updateable, BulkDeletable
 from labelbox.orm.model import Entity, Field, Relationship
-from labelbox.schema.asset_attachment import AssetAttachment
+
+if TYPE_CHECKING:
+    from labelbox import AssetAttachment
 
 logger = logging.getLogger(__name__)
 
@@ -41,7 +43,7 @@ class DataRow(DbObject, Updateable, BulkDeletable):
     attachments = Relationship.ToMany("AssetAttachment", False, "attachments")
 
     supported_meta_types = supported_attachment_types = set(
-        AssetAttachment.AttachmentType.__members__)
+        Entity.AssetAttachment.AttachmentType.__members__)
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -58,7 +60,7 @@ def bulk_delete(data_rows) -> None:
         BulkDeletable._bulk_delete(data_rows, True)
 
     def create_attachment(self, attachment_type,
-                          attachment_value) -> AssetAttachment:
+                          attachment_value) -> "AssetAttachment":
         """ Adds an AssetAttachment to a DataRow.
             Labelers can view these attachments while labeling.
 
@@ -73,7 +75,7 @@ def create_attachment(self, attachment_type,
         Raises:
             ValueError: asset_type must be one of the supported types.
         """
-        AssetAttachment.validate_attachment_type(attachment_type)
+        Entity.AssetAttachment.validate_attachment_type(attachment_type)
 
         attachment_type_param = "type"
         attachment_value_param = "value"

labelbox/schema/dataset.py

@@ -1,4 +1,4 @@
-from typing import Generator, List, Union, Any
+from typing import Generator, List, Union, Any, TYPE_CHECKING
 import os
 import json
 import logging
@@ -7,19 +7,18 @@
 import ndjson
 from itertools import islice
 
-from labelbox.data.serialization.ndjson.base import DataRow
-from labelbox.schema import iam_integration
-from labelbox.schema.task import Task
-from labelbox.schema.user import User
-from labelbox import utils
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import StringIO
 import requests
 
+from labelbox import utils
 from labelbox.exceptions import InvalidQueryError, LabelboxError, ResourceNotFoundError, InvalidAttributeError
 from labelbox.orm.db_object import DbObject, Updateable, Deletable
 from labelbox.orm.model import Entity, Field, Relationship
 
+if TYPE_CHECKING:
+    from labelbox import Task, User, DataRow
+
 logger = logging.getLogger(__name__)
 
 
@@ -53,7 +52,7 @@ class Dataset(DbObject, Updateable, Deletable):
     iam_integration = Relationship.ToOne("IAMIntegration", False,
                                          "iam_integration", "signer")
 
-    def create_data_row(self, **kwargs) -> DataRow:
+    def create_data_row(self, **kwargs) -> "DataRow":
         """ Creates a single DataRow belonging to this dataset.
 
         >>> dataset.create_data_row(row_data="http://my_site.com/photos/img_01.jpg")
@@ -122,7 +121,7 @@ def create_data_rows_sync(self, items) -> None:
             url_param: descriptor_url
         })
 
-    def create_data_rows(self, items) -> Union[Task, List[Any]]:
+    def create_data_rows(self, items) -> Union["Task", List[Any]]:
         """ Asynchronously bulk upload data rows
 
         Use this instead of `Dataset.create_data_rows_sync` uploads for batches that contain more than 1000 data rows.
@@ -317,7 +316,9 @@ def convert_item(item):
         data = json.dumps(items)
         return self.client.upload_data(data)
 
-    def data_rows_for_external_id(self, external_id, limit=10) -> List[DataRow]:
+    def data_rows_for_external_id(self,
+                                  external_id,
+                                  limit=10) -> List["DataRow"]:
         """ Convenience method for getting a multiple `DataRow` belonging to this
         `Dataset` that has the given `external_id`.
 
@@ -344,7 +345,7 @@ def data_rows_for_external_id(self, external_id, limit=10) -> List[DataRow]:
             raise ResourceNotFoundError(DataRow, where)
         return data_rows
 
-    def data_row_for_external_id(self, external_id) -> DataRow:
+    def data_row_for_external_id(self, external_id) -> "DataRow":
         """ Convenience method for getting a single `DataRow` belonging to this
         `Dataset` that has the given `external_id`.
 

labelbox/schema/label.py

@@ -1,8 +1,11 @@
-from labelbox.data.serialization.labelbox_v1.label import Review
+from typing import TYPE_CHECKING
+
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Updateable, BulkDeletable
 from labelbox.orm.model import Entity, Field, Relationship
-from labelbox.schema.benchmark import Benchmark
+
+if TYPE_CHECKING:
+    from labelbox import Benchmark, Review
 """ Client-side object type definitions. """
 
 
@@ -50,7 +53,7 @@ def bulk_delete(labels) -> None:
         """
         BulkDeletable._bulk_delete(labels, False)
 
-    def create_review(self, **kwargs) -> Review:
+    def create_review(self, **kwargs) -> "Review":
         """ Creates a Review for this label.
 
         Args:
@@ -60,7 +63,7 @@ def create_review(self, **kwargs) -> Review:
         kwargs[Entity.Review.project.name] = self.project()
         return self.client._create(Entity.Review, kwargs)
 
-    def create_benchmark(self) -> Benchmark:
+    def create_benchmark(self) -> "Benchmark":
         """ Creates a Benchmark for this Label.
 
         Returns:

labelbox/schema/model.py

@@ -1,7 +1,10 @@
+from typing import TYPE_CHECKING
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Entity, Field, Relationship
-from labelbox.schema.model_run import ModelRun
+
+if TYPE_CHECKING:
+    from labelbox import ModelRun
 
 
 class Model(DbObject):
@@ -15,7 +18,7 @@ class Model(DbObject):
     name = Field.String("name")
     model_runs = Relationship.ToMany("ModelRun", False)
 
-    def create_model_run(self, name) -> ModelRun:
+    def create_model_run(self, name) -> "ModelRun":
         """ Creates a model run belonging to this model.
 
         Args:

labelbox/schema/model_run.py

@@ -1,14 +1,16 @@
-from typing import Dict, Iterable, Union
+from typing import TYPE_CHECKING, Dict, Iterable, Union
 from pathlib import Path
 import os
 import time
 
 from labelbox.pagination import PaginatedCollection
-from labelbox.schema.annotation_import import MEAPredictionImport
 from labelbox.orm.query import results_query_part
-from labelbox.orm.model import Field, Relationship
+from labelbox.orm.model import Field, Relationship, Entity
 from labelbox.orm.db_object import DbObject
 
+if TYPE_CHECKING:
+    from labelbox import MEAPredictionImport
+
 
 class ModelRun(DbObject):
     name = Field.String("name")
@@ -119,13 +121,13 @@ def add_predictions(
         kwargs = dict(client=self.client, model_run_id=self.uid, name=name)
         if isinstance(predictions, str) or isinstance(predictions, Path):
             if os.path.exists(predictions):
-                return MEAPredictionImport.create_from_file(
+                return Entity.MEAPredictionImport.create_from_file(
                     path=str(predictions), **kwargs)
             else:
-                return MEAPredictionImport.create_from_url(url=str(predictions),
-                                                           **kwargs)
+                return Entity.MEAPredictionImport.create_from_url(
+                    url=str(predictions), **kwargs)
         elif isinstance(predictions, Iterable):
-            return MEAPredictionImport.create_from_objects(
+            return Entity.MEAPredictionImport.create_from_objects(
                 predictions=predictions, **kwargs)
         else:
             raise ValueError(

labelbox/schema/ontology.py

@@ -7,7 +7,6 @@
 
 from pydantic import constr
 
-# from labelbox.schema import project
 from labelbox.exceptions import InconsistentOntologyException
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field, Relationship

labelbox/schema/organization.py

@@ -1,13 +1,12 @@
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional
 
 from labelbox.exceptions import LabelboxError
 from labelbox import utils
 from labelbox.orm.db_object import DbObject, query, Entity
 from labelbox.orm.model import Field, Relationship
-from labelbox.schema.iam_integration import IAMIntegration
-from labelbox.schema.invite import Invite, InviteLimit, ProjectRole
-from labelbox.schema.user import User
-from labelbox.schema.role import Role
+
+if TYPE_CHECKING:
+    from labelbox import Role, User, ProjectRole, Invite, InviteLimit, IAMIntegration
 
 
 class Organization(DbObject):
@@ -46,8 +45,8 @@ def __init__(self, *args, **kwargs):
     def invite_user(
             self,
             email: str,
-            role: Role,
-            project_roles: Optional[List[ProjectRole]] = None) -> Invite:
+            role: "Role",
+            project_roles: Optional[List["ProjectRole"]] = None) -> "Invite":
         """
         Invite a new member to the org. This will send the user an email invite
 
@@ -73,7 +72,7 @@ def invite_user(
         data_param = "data"
         query_str = """mutation createInvitesPyApi($%s: [CreateInviteInput!]){
                     createInvites(data: $%s){  invite { id createdAt organizationRoleName inviteeEmail inviter { %s } }}}""" % (
-            data_param, data_param, query.results_query_part(User))
+            data_param, data_param, query.results_query_part(Entity.User))
 
         projects = [{
             "projectId": project_role.project.uid,
@@ -93,9 +92,9 @@ def invite_user(
         invite_response = res['createInvites'][0]['invite']
         if not invite_response:
             raise LabelboxError(f"Unable to send invite for email {email}")
-        return Invite(self.client, invite_response)
+        return Entity.Invite(self.client, invite_response)
 
-    def invite_limit(self) -> InviteLimit:
+    def invite_limit(self) -> "InviteLimit":
         """ Retrieve invite limits for the org
         This already accounts for users currently in the org
         Meaining that  `used = users + invites, remaining = limit - (users + invites)`
@@ -109,10 +108,10 @@ def invite_limit(self) -> InviteLimit:
             """query InvitesLimitPyApi($%s: ID!) {
             invitesLimit(where: {id: $%s}) { used limit remaining }
         }""" % (org_id_param, org_id_param), {org_id_param: self.uid})
-        return InviteLimit(
+        return Entity.InviteLimit(
             **{utils.snake_case(k): v for k, v in res['invitesLimit'].items()})
 
-    def remove_user(self, user: User) -> None:
+    def remove_user(self, user: "User") -> None:
         """
         Deletes a user from the organization. This cannot be undone without sending another invite.
 
@@ -126,7 +125,7 @@ def remove_user(self, user: User) -> None:
             updateUser(where: {id: $%s}, data: {deleted: true}) { id deleted }
         }""" % (user_id_param, user_id_param), {user_id_param: user.uid})
 
-    def get_iam_integrations(self) -> List[IAMIntegration]:
+    def get_iam_integrations(self) -> List["IAMIntegration"]:
         """
         Returns all IAM Integrations for an organization
         """
@@ -145,7 +144,7 @@ def get_iam_integrations(self) -> List[IAMIntegration]:
             for integration_data in res['iamIntegrations']
         ]
 
-    def get_default_iam_integration(self) -> Optional[IAMIntegration]:
+    def get_default_iam_integration(self) -> Optional["IAMIntegration"]:
         """
         Returns the default IAM integration for the organization.
         Will return None if there are no default integrations for the org.