Add Project.extend_reservations.

Author: florijanstamenkovic

labelbox/schema/project.py

@@ -241,6 +241,25 @@ def upsert_review_queue(self, quota_factor):
             query_str, {project_param: self.uid, quota_param: quota_factor})
 
 
+    def extend_reservations(self, queue_type):
+        """ Extend all the current reservations for the current user on the given
+        queue type.
+        Args:
+            queue_type (str): Either "LabelingQueue" or "ReviewQueue"
+        Return:
+            int, the number of reservations that were extended.
+        """
+        if queue_type not in ("LabelingQueue", "ReviewQueue"):
+            raise InvalidQueryError("Unsupported queue type: %s" % queue_type)
+
+        project_param = "projectId"
+        query_str = """mutation ExtendReservationsPyApi($%s: ID!){
+            extendReservations(projectId:$%s queueType:%s)}""" % (
+                project_param, project_param, queue_type)
+        res = self.client.execute(query_str, {project_param: self.uid})
+        return res["data"]["extendReservations"]
+
+
 class LabelingParameterOverride(DbObject):
     priority = Field.Int("priority")
     number_of_labels = Field.Int("number_of_labels")

tests/integration/test_project.py

@@ -1,4 +1,7 @@
+import pytest
+
 from labelbox import Project
+from labelbox.exceptions import InvalidQueryError
 
 
 def test_project(client, rand_gen):
@@ -56,3 +59,12 @@ def test_upsert_review_queue(client, rand_gen):
     project = client.create_project(name=rand_gen(str))
     project.upsert_review_queue(0.6)
     project.delete()
+
+
+def test_extend_reservations(client, rand_gen):
+    project = client.create_project(name=rand_gen(str))
+    assert project.extend_reservations("LabelingQueue") == 0
+    assert project.extend_reservations("ReviewQueue") == 0
+    with pytest.raises(InvalidQueryError):
+        project.extend_reservations("InvalidQueueType")
+    project.delete()