recommended changes

Author: Matt Sokoloff

labelbox/client.py

@@ -124,11 +124,8 @@ def convert_value(value):
         data = json.dumps({'query': query, 'variables': params}).encode('utf-8')
 
         try:
-            endpoint = self.endpoint
-            if beta:
-                endpoint = endpoint.replace('/graphql', '/_gql')
-
-            response = requests.post(endpoint,
+            response = requests.post(self.endpoint.replace('/graphql', '/_gql')
+                                     if beta else self.endpoint,
                                      data=data,
                                      headers=self.headers,
                                      timeout=timeout)

labelbox/orm/db_object.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from functools import partial
+from functools import wraps
 import logging
 
 from labelbox import utils
@@ -258,11 +258,12 @@ def delete(self):
 
 def beta(fn):
 
+    @wraps(fn)
     def wrapper(self, *args, **kwargs):
         if not self.client.enable_beta:
             raise Exception(
                 f"This function {fn.__name__} relies on a beta feature in the api. This means that the interface could change."
-                " Set `enable_beta=True` in the client to enable use of these functions."
+                " Set `enable_beta=True` in the client to enable use of beta functions."
             )
         return fn(self, *args, **kwargs)
 

labelbox/pagination.py

@@ -1,8 +1,11 @@
 # Size of a single page in a paginated query.
 from abc import ABC, abstractmethod
-from labelbox.orm.db_object import DbObject
-from typing import Any, Dict, List, Optional
-from labelbox import Client
+from typing import Any, Callable, Dict, List, Optional
+
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from labelbox import Client
+    from labelbox.orm.db_object import DbObject
 
 _PAGE_SIZE = 100
 
@@ -17,11 +20,11 @@ class PaginatedCollection:
     """
 
     def __init__(self,
-                 client: Client,
+                 client: "Client",
                  query: str,
                  params: Dict[str, str],
                  dereferencing: Dict[str, Any],
-                 obj_class: DbObject,
+                 obj_class: "DbObject",
                  cursor_path: Optional[Dict[str, Any]] = None,
                  beta: bool = False):
         """ Creates a PaginatedCollection.
@@ -45,12 +48,14 @@ def __init__(self,
         self.params = params
         self.dereferencing = dereferencing
         self.obj_class = obj_class
+        self.cursor_path = cursor_path
         self.beta = beta
 
         self._fetched_all = False
         self._data: List[Dict[str, Any]] = []
         self._data_ind = 0
-        self.cursor_path = _CursorPagination(
+
+        self.paginator = _CursorPagination(
             client, cursor_path) if cursor_path else _OffsetPagination(client)
 
     def __iter__(self):
@@ -97,7 +102,7 @@ def fetch_results(self, query: str, params: Dict[str, Any],
 
 class _CursorPagination(_Pagination):
 
-    def __init__(self, client: Client, cursor_path: Dict[str, Any]):
+    def __init__(self, client: "Client", cursor_path: Dict[str, Any]):
         self.client = client
         self.cursor_path = cursor_path
         self.next_cursor: Optional[str] = None
@@ -119,7 +124,7 @@ def fetch_results(self, query: str, params: Dict[str, Any],
 
 class _OffsetPagination(_Pagination):
 
-    def __init__(self, client):
+    def __init__(self, client: "Client"):
         self.client = client
         self._fetched_pages = 0