try/finally beta

Author: Matt Sokoloff

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.7
+FROM python:3.6
 
 RUN pip install pytest pytest-cases
 

labelbox/client.py

@@ -3,7 +3,6 @@
 import logging
 import mimetypes
 import os
-from typing import Tuple
 
 from google.api_core import retry
 import requests
@@ -63,24 +62,25 @@ def __init__(self,
                     "Labelbox API key not provided")
             api_key = os.environ[_LABELBOX_API_KEY]
         self.api_key = api_key
-        self.enable_beta = enable_beta
-
+    
+        self.enable_beta = True
         if enable_beta:
             logger.info("Beta features have been enabled")
+        
 
         logger.info("Initializing Labelbox client at '%s'", endpoint)
-
         self.endpoint = endpoint
         self.headers = {
             'Accept': 'application/json',
             'Content-Type': 'application/json',
             'Authorization': 'Bearer %s' % api_key,
             'X-User-Agent': f'python-sdk {SDK_VERSION}'
         }
+    
 
     @retry.Retry(predicate=retry.if_exception_type(
         labelbox.exceptions.InternalServerError))
-    def execute(self, query, params=None, timeout=30.0):
+    def execute(self, query, params=None, timeout=30.0, beta = False):
         """ Sends a request to the server for the execution of the
         given query.
 
@@ -126,7 +126,11 @@ def convert_value(value):
         data = json.dumps({'query': query, 'variables': params}).encode('utf-8')
 
         try:
-            response = requests.post(self.endpoint,
+            endpoint = self.endpoint
+            if beta:
+                endpoint = endpoint.replace('/graphql', '/_gql')
+
+            response = requests.post(endpoint,
                                      data=data,
                                      headers=self.headers,
                                      timeout=timeout)

labelbox/orm/db_object.py

@@ -1,6 +1,8 @@
 from dataclasses import field
 from datetime import datetime, timezone
+from functools import partial
 import logging
+from os import EX_CANTCREAT
 
 from labelbox import utils
 from labelbox.exceptions import InvalidQueryError, InvalidAttributeError
@@ -45,9 +47,10 @@ def __init__(self, client, field_values):
         self._set_field_values(field_values)
 
         for relationship in self.relationships():
-            value = field_values[
-                relationship.name] if relationship.precompute else None
-
+            value =   field_values.get(relationship.name)
+            if relationship.precompute and value is None:
+                raise KeyError(f"Expected field  values for {relationship.name}")
+            
             setattr(self, relationship.name,
                     RelationshipManager(self, relationship, value))
 
@@ -254,18 +257,23 @@ def delete(self):
         type(self).bulk_delete([self])
 
 
-
 def beta(fn):
     def wrapper(self, *args, **kwargs):
         if not isinstance(self, DbObject):
             raise TypeError("Cannot decorate functions that are not functions of `DbOjects` with `beta` decorator")
         if not self.client.enable_beta:
-            logger.warning(
+            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 silence this warning.")
-            self.client.endpoint = self.client.endpoint.replace("/graphql", "/_gql")
+                " Set `enable_beta=True` in the client to enable use of these functions.")
+        execute_fn = self.client.execute            
+        try:
+            self.client.execute = partial(execute_fn, beta = True)
             result = fn(self, *args, **kwargs)
-            self.client.endpoint = self.client.endpoint.replace("/_gql", "/graphql")
+            if isinstance(result, PaginatedCollection):
+                result.beta = True
             return result
+        finally:
+            self.client.execute = execute_fn
     return wrapper
 
+

labelbox/pagination.py

@@ -17,8 +17,8 @@ def __init__(self,
                  params,
                  dereferencing,
                  obj_class,
-                 cursor_path=None,
-                 beta=False):
+                 cursor_path=None
+                 ):
         """ Creates a PaginatedCollection.
 
         Args:
@@ -40,7 +40,7 @@ def __init__(self,
         self.params = params
         self.dereferencing = dereferencing
         self.obj_class = obj_class
-        self.beta = beta
+        self.beta = False
 
         self._fetched_all = False
         self._data = []
@@ -95,13 +95,11 @@ def get_next_cursor(self, 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
+        return self.next_cursor is None
 
-    def fetch_results(self, query, params, experimental):
+    def fetch_results(self, query, params, beta):
         params.update({'from': self.next_cursor, 'first': _PAGE_SIZE})
-        return self.client.execute(query, params, experimental=experimental)
+        return self.client.execute(query, params, beta=beta)
 
 
 class _OffsetPagination:
@@ -116,6 +114,6 @@ def fetched_all(self, n_items, results):
             return True
         return False
 
-    def fetch_results(self, query, params, experimental):
+    def fetch_results(self, query, params, beta):
         query = query % (self._fetched_pages * _PAGE_SIZE, _PAGE_SIZE)
-        return self.client.execute(query, params, experimental=experimental)
+        return self.client.execute(query, params, beta=beta)

labelbox/schema/invite.py

@@ -58,5 +58,4 @@ def revoke(self):
                cancelInvite(where: $where) {id}}"""
         self.client.execute(query_str, {'where': {
             'id': self.uid
-        }},
-                            experimental=True)
+        }})

labelbox/schema/organization.py

@@ -60,8 +60,7 @@ def invites(self) -> PaginatedCollection:
             self.client,
             query_str, {}, ['organization', 'invites', 'nodes'],
             Invite,
-            cursor_path=['organization', 'invites', 'nextCursor'],
-            experimental=True)
+            cursor_path=['organization', 'invites', 'nextCursor'])
 
     def _assign_user_role(self, email: str, role: Role,
                           project_roles: List[ProjectRole]) -> Dict[str, Any]:
@@ -95,7 +94,8 @@ def _assign_user_role(self, email: str, role: Role,
                     "projects": projects
                 }]
             },
-            experimental=True)  # We prob want to return an invite
+            ) 
+        # We prob want to return an invite
         # Could support bulk ops in the future
         invite_info = res['createInvites'][0]['invite']
         return invite_info
@@ -159,7 +159,7 @@ def user_limit(self) -> UserLimit:
         query_str = """query UsersLimitPyApi { 
             organization {id account { id usersLimit { dateLimitWasReached remaining used limit }}}}
         """
-        res = self.client.execute(query_str, experimental=True)
+        res = self.client.execute(query_str)
         return UserLimit(
             **{
                 utils.snake_case(k): v for k, v in res['organization']
@@ -179,12 +179,11 @@ def invite_limit(self) -> InviteLimit:
         org_id_param = "organizationId"
         res = self.client.execute("""query InvitesLimitPyApi($%s: ID!) {
             invitesLimit(where: {id: $%s}) { used limit remaining }
-        }""" % (org_id_param, org_id_param), {org_id_param: self.uid},
-                                  experimental=True)
+        }""" % (org_id_param, org_id_param), {org_id_param: self.uid})
         return InviteLimit(
             **{utils.snake_case(k): v for k, v in res['invitesLimit'].items()})
 
-
+    @beta
     def remove_user(self, user: User):
         """
         Deletes a user from the organization. This cannot be undone without sending another invite.