update pagination, other small changes

Author: Matt Sokoloff

labelbox/pagination.py

@@ -42,12 +42,14 @@ def __init__(self,
         self.obj_class = obj_class
         self.experimental = experimental
 
-        self._fetched_pages = 0
         self._fetched_all = False
         self._data = []
         self._data_ind = 0
-        self.next_cursor = None
-        self.cursor_path = cursor_path
+
+        if cursor_path:
+            self.paginator = _CursorPagination(client, cursor_path)
+        else:
+            self.paginator = _OffsetPagination(client)
 
     def __iter__(self):
         self._data_ind = 0
@@ -56,47 +58,64 @@ def __iter__(self):
     def get_page_data(self, results):
         for deref in self.dereferencing:
             results = results[deref]
-
         return [self.obj_class(self.client, result) for result in results]
 
-    def get_next_cursor(self, results):
-        for path in self.cursor_path:
-            results = results[path]
-        return results
-
-    def get_query(self):
-        if self.cursor_path:
-            # Cursor Pagination
-            self.params.update({'from': self.next_cursor, 'first': _PAGE_SIZE})
-            return self.query
-
-        # Offset Pagination
-        return self.query % (self._fetched_pages * _PAGE_SIZE, _PAGE_SIZE)
-
     def __next__(self):
         if len(self._data) <= self._data_ind:
             if self._fetched_all:
                 raise StopIteration()
 
-            results = self.client.execute(self.get_query(),
-                                          self.params,
-                                          experimental=self.experimental)
-            self._fetched_pages += 1
-
+            results = self.paginator.fetch_results(self.query, self.params,
+                                                   self.experimental)
             page_data = self.get_page_data(results)
             self._data.extend(page_data)
+            n_items = len(page_data)
 
-            if len(page_data) < _PAGE_SIZE:
-                self._fetched_all = True
-
-            if len(page_data) == 0:
+            if n_items == 0:
                 raise StopIteration()
 
-            if self.cursor_path:
-                self.next_cursor = self.get_next_cursor(results)
-                if self.next_cursor is None:
-                    self._fetched_all = True
+            self._fetched_all = self.paginator.fetched_all(n_items, results)
 
         rval = self._data[self._data_ind]
         self._data_ind += 1
         return rval
+
+
+class _CursorPagination:
+
+    def __init__(self, client, cursor_path):
+        self.client = client
+        self.cursor_path = cursor_path
+        self.next_cursor = None
+
+    def get_next_cursor(self, results):
+        for path in self.cursor_path:
+            results = results[path]
+        return results
+
+    def fetched_all(self, n_items, results):
+        self.next_cursor = self.get_next_cursor(results)
+        if self.next_cursor is None:
+            return True
+        return False
+
+    def fetch_results(self, query, params, experimental):
+        params.update({'from': self.next_cursor, 'first': _PAGE_SIZE})
+        return self.client.execute(query, params, experimental=experimental)
+
+
+class _OffsetPagination:
+
+    def __init__(self, client):
+        self.client = client
+        self._fetched_pages = 0
+
+    def fetched_all(self, n_items, results):
+        self._fetched_pages += 1
+        if n_items < _PAGE_SIZE:
+            return True
+        return False
+
+    def fetch_results(self, query, params, experimental):
+        query = query % (self._fetched_pages * _PAGE_SIZE, _PAGE_SIZE)
+        return self.client.execute(query, params, experimental=experimental)

labelbox/schema/invite.py

@@ -45,10 +45,9 @@ def __init__(self, client, invite_response):
         super().__init__(client, invite_response)
 
         self.project_roles = [
-            ProjectRole(
-                project=client.get_project(r['projectId']),
-                role=client.get_roles()[r['projectRoleName'].upper().replace(
-                    " ", "_")]) for r in project_roles
+            ProjectRole(project=client.get_project(r['projectId']),
+                        role=client.get_roles()[r['projectRoleName']])
+            for r in project_roles
         ]
 
     def revoke(self):

labelbox/schema/organization.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from labelbox.exceptions import LabelboxError
 from labelbox import utils
@@ -78,9 +78,9 @@ def _assign_user_role(self, email: str, role: Role,
             data_param, data_param)
 
         projects = [{
-            "projectId": x.project.uid,
-            "projectRoleId": x.role.uid
-        } for x in project_roles]
+            "projectId": project_role.project.uid,
+            "projectRoleId": project_role.role.uid
+        } for project_role in project_roles]
 
         res = self.client.execute(
             query_str, {
@@ -97,30 +97,29 @@ def _assign_user_role(self, email: str, role: Role,
         invite_info = res['createInvites'][0]['invite']
         return invite_info
 
-    def invite_user(self,
-                    email: str,
-                    role: Role,
-                    project_roles: List[ProjectRole] = []) -> Invite:
+    def invite_user(
+            self,
+            email: str,
+            role: Role,
+            project_roles: Optional[List[ProjectRole]] = None) -> Invite:
         """
         Invite a new member to the org. This will send the user an email invite
 
         Args:
             email (str): email address of the user to invite
             role (Role): Role to assign to the user
-            project_roles (List[ProjectRoles]): List of project roles to assign to the User (if they have a project based org role).
+            project_roles (Optional[List[ProjectRoles]]): List of project roles to assign to the User (if they have a project based org role).
 
         Returns:
             Invite for the user
 
         """
         remaining_invites = self.invite_limit().remaining
         if remaining_invites == 0:
-            if next(self.users((User.email == email) &
-                               (User.deleted == True)), None) is None:
-                raise LabelboxError(
-                    "Invite(s) cannot be sent because you do not have enough available seats in your organization. "
-                    "Please upgrade your account, revoke pending invitations or remove other users."
-                )
+            raise LabelboxError(
+                "Invite(s) cannot be sent because you do not have enough available seats in your organization. "
+                "Please upgrade your account, revoke pending invitations or remove other users."
+            )
         for invite in self.invites():
             if invite.email == email:
                 raise ValueError(
@@ -130,18 +129,19 @@ def invite_user(self,
         if not isinstance(role, Role):
             raise TypeError(f"role must be Role type. Found {role}")
 
-        if len(project_roles) and role.name != "NONE":
+        if project_roles and role.name != "NONE":
             raise ValueError(
                 "Project roles cannot be set for a user with organization level permissions. Found role name `{role.name}`, expected `NONE`"
             )
 
-        for project_role in project_roles:
+        _project_roles = [] if project_roles is None else project_roles
+        for project_role in _project_roles:
             if not isinstance(project_role, ProjectRole):
                 raise TypeError(
                     f"project_roles must be a list of `ProjectRole`s. Found {project_role}"
                 )
 
-        invite_response = self._assign_user_role(email, role, project_roles)
+        invite_response = self._assign_user_role(email, role, _project_roles)
         return Invite(self.client, invite_response)
 
     def user_limit(self) -> UserLimit:

labelbox/schema/role.py

@@ -38,12 +38,12 @@ def __repr__(self):
         return str({k: getattr(self, k) for k in self.valid_roles})
 
     def __getitem__(self, name):
-        if name in self.valid_roles:
-            return getattr(self, name)
-        else:
+        name = name.replace(' ', '_').upper()
+        if name not in self.valid_roles:
             raise ValueError(
                 f"No role named {name} exists. Valid names are one of {self.valid_roles}"
             )
+        return getattr(self, name)
 
     def __iter__(self):
         self.key_iter = iter(self.valid_roles)