updates to type hinting and adding to some inits as well as providing additional errors

Author: jtsodapop

labelbox/__init__.py

@@ -13,11 +13,13 @@
 from labelbox.schema.user import User
 from labelbox.schema.organization import Organization
 from labelbox.schema.task import Task
-from labelbox.schema.labeling_frontend import LabelingFrontend
+from labelbox.schema.labeling_frontend import LabelingFrontend, LabelingFrontendOptions
 from labelbox.schema.asset_attachment import AssetAttachment
 from labelbox.schema.webhook import Webhook
 from labelbox.schema.ontology import Ontology, OntologyBuilder, Classification, Option, Tool, FeatureSchema
 from labelbox.schema.role import Role, ProjectRole
 from labelbox.schema.invite import Invite, InviteLimit
 from labelbox.schema.data_row_metadata import DataRowMetadataOntology
 from labelbox.schema.model_run import ModelRun
+from labelbox.schema.benchmark import Benchmark
+from labelbox.schema.iam_integration import IAMIntegration

labelbox/client.py

@@ -29,6 +29,9 @@
 from labelbox.schema.model import Model
 from labelbox.schema.ontology import Ontology, Tool, Classification
 from labelbox.schema.organization import Organization
+from labelbox.schema.user import User
+from labelbox.schema.project import Project
+from labelbox.schema.role import Role
 
 logger = logging.getLogger(__name__)
 

labelbox/data/serialization/labelbox_v1/converter.py

@@ -1,5 +1,5 @@
 from labelbox.data.serialization.labelbox_v1.objects import LBV1Mask
-from typing import Any, Dict, Generator, Iterable
+from typing import Any, Dict, Generator, Iterable, Union
 import logging
 
 import ndjson
@@ -19,7 +19,7 @@
 class LBV1Converter:
 
     @staticmethod
-    def deserialize_video(json_data: Iterable[Dict[str, Any]],
+    def deserialize_video(json_data: Union[str, Iterable[Dict[str, Any]]],
                           client: "labelbox.Client") -> LabelGenerator:
         """
         Converts a labelbox video export into the common labelbox format.
@@ -36,7 +36,8 @@ def deserialize_video(json_data: Iterable[Dict[str, Any]],
         return LabelGenerator(data=label_generator)
 
     @staticmethod
-    def deserialize(json_data: Iterable[Dict[str, Any]]) -> LabelGenerator:
+    def deserialize(
+            json_data: Union[str, Iterable[Dict[str, Any]]]) -> LabelGenerator:
         """
         Converts a labelbox export (non-video) into the common labelbox format.
 

labelbox/orm/model.py

@@ -1,6 +1,7 @@
 from enum import Enum, auto
-from typing import Dict, List, Union
+from typing import Dict, List, Union, Any, Type, TYPE_CHECKING
 
+import labelbox
 from labelbox import utils
 from labelbox.exceptions import InvalidAttributeError
 from labelbox.orm.comparison import Comparison
@@ -245,6 +246,9 @@ class EntityMeta(type):
     # Maps Entity name to Relationships for all currently defined Entities
     relationship_mappings: Dict[str, List[Relationship]] = {}
 
+    def __setattr__(self, key: Any, value: Any):
+        super().__setattr__(key, value)
+
     def __init__(cls, clsname, superclasses, attributedict):
         super().__init__(clsname, superclasses, attributedict)
         cls.validate_cached_relationships()
@@ -325,6 +329,21 @@ class Entity(metaclass=EntityMeta):
     # suchs as `fields()`.
     deleted = Field.Boolean("deleted")
 
+    if TYPE_CHECKING:
+        DataRow: Type[labelbox.DataRow]
+        Webhook: Type[labelbox.Webhook]
+        Task: Type[labelbox.Task]
+        AssetAttachment: Type[labelbox.AssetAttachment]
+        ModelRun: Type[labelbox.ModelRun]
+        Review: Type[labelbox.Review]
+        User: Type[labelbox.User]
+        LabelingFrontend: Type[labelbox.LabelingFrontend]
+        BulkImportRequest: Type[labelbox.BulkImportRequest]
+        Benchmark: Type[labelbox.Benchmark]
+        IAMIntegration: Type[labelbox.IAMIntegration]
+        LabelingFrontendOptions: Type[labelbox.LabelingFrontendOptions]
+        Label: Type[labelbox.Label]
+
     @classmethod
     def _attributes_of_type(cls, attr_type):
         """ Yields all the attributes in `cls` of the given `attr_type`. """

labelbox/pagination.py

@@ -1,6 +1,6 @@
 # Size of a single page in a paginated query.
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, List, Optional, Tuple, Type
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -23,8 +23,8 @@ def __init__(self,
                  client: "Client",
                  query: str,
                  params: Dict[str, str],
-                 dereferencing: Dict[str, Any],
-                 obj_class: Type["DbObject"],
+                 dereferencing: Union[List[str], Dict[str, Any]],
+                 obj_class: Union[Type["DbObject"], Callable[[Any, Any], Any]],
                  cursor_path: Optional[Dict[str, Any]] = None,
                  experimental: bool = False):
         """ Creates a PaginatedCollection.

labelbox/schema/dataset.py

@@ -1,4 +1,4 @@
-from typing import Generator, List
+from typing import Generator, List, Union, Any
 import os
 import json
 import logging
@@ -10,6 +10,7 @@
 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
@@ -121,7 +122,7 @@ def create_data_rows_sync(self, items) -> None:
             url_param: descriptor_url
         })
 
-    def create_data_rows(self, items) -> Task:
+    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.
@@ -164,14 +165,15 @@ def create_data_rows(self, items) -> Task:
 
         # Fetch and return the task.
         task_id = res["taskId"]
-        user = self.client.get_user()
-        task = list(user.created_tasks(where=Entity.Task.uid == task_id))
+        user: User = self.client.get_user()
+        tasks: List[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:
+        if len(tasks) != 1:
             raise ResourceNotFoundError(Entity.Task, task_id)
-        task = task[0]
+        task: Task = tasks[0]
         task._user = user
         return task
 

labelbox/schema/ontology.py

@@ -7,7 +7,7 @@
 
 from pydantic import constr
 
-from labelbox.schema import project
+# 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/project.py

@@ -6,7 +6,7 @@
 from collections import namedtuple
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Dict, Union, Iterable, List, Optional
+from typing import Dict, Union, Iterable, List, Optional, Any
 from urllib.parse import urlparse
 
 import ndjson
@@ -203,6 +203,10 @@ def video_label_generator(self, timeout_seconds=600) -> LabelGenerator:
         _check_converter_import()
         json_data = self.export_labels(download=True,
                                        timeout_seconds=timeout_seconds)
+        # assert that the instance this would fail is only if timeout runs out
+        assert isinstance(
+            json_data,
+            List), "Unable to successfully get labels. Please try again"
         if json_data is None:
             raise TimeoutError(
                 f"Unable to download labels in {timeout_seconds} seconds."
@@ -227,6 +231,10 @@ def label_generator(self, timeout_seconds=600) -> LabelGenerator:
         _check_converter_import()
         json_data = self.export_labels(download=True,
                                        timeout_seconds=timeout_seconds)
+        # assert that the instance this would fail is only if timeout runs out
+        assert isinstance(
+            json_data,
+            List), "Unable to successfully get labels. Please try again"
         if json_data is None:
             raise TimeoutError(
                 f"Unable to download labels in {timeout_seconds} seconds."
@@ -241,7 +249,10 @@ def label_generator(self, timeout_seconds=600) -> LabelGenerator:
                 "Or use project.video_label_generator() for video data.")
         return LBV1Converter.deserialize(json_data)
 
-    def export_labels(self, download=False, timeout_seconds=600) -> str:
+    def export_labels(
+            self,
+            download=False,
+            timeout_seconds=600) -> Optional[Union[str, List[Dict[Any, Any]]]]:
         """ Calls the server-side Label exporting that generates a JSON
         payload, and returns the URL to that payload.
 
@@ -564,7 +575,7 @@ def _update_queue_mode(self, mode: QueueMode) -> QueueMode:
 
         return mode
 
-    def queue_mode(self) -> str:
+    def queue_mode(self) -> QueueMode:
         """Provides the status of if queue mode is enabled in the project."""
 
         query_str = """query %s($projectId: ID!) {

labelbox/schema/task.py

@@ -1,9 +1,11 @@
 import logging
 import time
+from typing import Optional
 
 from labelbox.exceptions import ResourceNotFoundError
 from labelbox.orm.db_object import DbObject
-from labelbox.orm.model import Field, Relationship
+from labelbox.orm.model import Field, Relationship, Entity
+from labelbox.schema.user import User
 
 logger = logging.getLogger(__name__)
 
@@ -27,13 +29,15 @@ class Task(DbObject):
     name = Field.String("name")
     status = Field.String("status")
     completion_percentage = Field.Float("completion_percentage")
+    _user: Optional[User] = None
 
     # Relationships
     created_by = Relationship.ToOne("User", False, "created_by")
     organization = Relationship.ToOne("Organization")
 
     def refresh(self) -> None:
         """ Refreshes Task data from the server. """
+        assert self._user is not None
         tasks = list(self._user.created_tasks(where=Task.uid == self.uid))
         if len(tasks) != 1:
             raise ResourceNotFoundError(Task, self.uid)