FIX: typing

Author: gdj0nes

labelbox/__init__.py

@@ -9,7 +9,7 @@
 from labelbox.schema.dataset import Dataset
 from labelbox.schema.data_row import DataRow
 from labelbox.schema.label import Label
-from labelbox.schema.batch import Batch, BatchPriority
+from labelbox.schema.batch import Batch
 from labelbox.schema.review import Review
 from labelbox.schema.user import User
 from labelbox.schema.organization import Organization

labelbox/orm/model.py

@@ -347,6 +347,7 @@ class Entity(metaclass=EntityMeta):
         Invite: Type[labelbox.Invite]
         InviteLimit: Type[labelbox.InviteLimit]
         ProjectRole: Type[labelbox.ProjectRole]
+        Batch: Type[labelbox.Batch]
 
     @classmethod
     def _attributes_of_type(cls, attr_type):

labelbox/schema/batch.py

@@ -1,10 +1,6 @@
-from typing import NewType
-
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field, Relationship
 
-BatchPriority = NewType('BatchPriority', int)
-
 
 class Batch(DbObject):
     """ A Batch is a group of data rows submitted to a project for labeling
@@ -26,5 +22,4 @@ class Batch(DbObject):
 
     # Relationships
     project = Relationship.ToOne("Project")
-
     created_by = Relationship.ToOne("User")

labelbox/schema/project.py

@@ -571,8 +571,8 @@ def setup(self, labeling_frontend, labeling_frontend_options) -> None:
 
     def create_batch(self,
                      name: str,
-                     data_rows: List[Union[str, "DataRow"]],
-                     priority: "BatchPriority" = 5):
+                     data_rows: List[str],
+                     priority: int = 5):
         """Create a new batch for a project
 
         Args:
@@ -588,8 +588,8 @@ def create_batch(self,
 
         dr_ids = []
         for dr in data_rows:
-            from labelbox import DataRow
-            if isinstance(dr, DataRow):
+
+            if isinstance(dr, Entity.DataRow):
                 dr_ids.append(dr.uid)
             elif isinstance(dr, str):
                 dr_ids.append(dr)