Merge pull request #19 from scaleapi/felix/batch

Add support for creating and finalizing task batch

Author: felixlaumon

scaleapi/__init__.py

@@ -2,6 +2,7 @@
 import requests
 
 from .tasks import Task
+from .batches import Batch
 
 TASK_TYPES = [
     'annotation',
@@ -10,15 +11,15 @@
     'comparison',
     'cuboidannotation',
     'datacollection',
-    'imageannotation', 
+    'imageannotation',
     'lineannotation',
     'namedentityrecognition',
     'pointannotation',
     'polygonannotation',
     'segmentannotation',
     'transcription',
-    'videoannotation', 
-    'videoboxannotation', 
+    'videoannotation',
+    'videoboxannotation',
     'videocuboidannotation'
 ]
 SCALE_ENDPOINT = 'https://api.scale.com/v1/'
@@ -35,27 +36,37 @@ class ScaleInvalidRequest(ScaleException, ValueError):
     pass
 
 
-class Tasklist(list):
+class Paginator(list):
     def __init__(self, docs, total, limit, offset, has_more, next_token=None):
-        super(Tasklist, self).__init__(docs)
+        super(Paginator, self).__init__(docs)
         self.docs = docs
         self.total = total
         self.limit = limit
         self.offset = offset
         self.has_more = has_more
         self.next_token = next_token
 
+
+class Tasklist(Paginator):
+    pass
+
+
+class Batchlist(Paginator):
+    pass
+
+
 class ScaleClient(object):
     def __init__(self, api_key):
         self.api_key = api_key
 
-    def _getrequest(self, endpoint, params={}):
+    def _getrequest(self, endpoint, params=None):
         """Makes a get request to an endpoint.
 
         If an error occurs, assumes that endpoint returns JSON as:
             { 'status_code': XXX,
               'error': 'I failed' }
         """
+        params = params or {}
         r = requests.get(SCALE_ENDPOINT + endpoint,
                          headers={"Content-Type": "application/json"},
                          auth=(self.api_key, ''), params=params)
@@ -114,7 +125,7 @@ def cancel_task(self, task_id):
     def tasks(self, **kwargs):
         """Returns a list of your tasks.
         Returns up to 100 at a time, to get more, use the next_token param passed back.
-        
+
         Note that offset is deprecated.
 
         start/end_time are ISO8601 dates, the time range of tasks to fetch.
@@ -125,7 +136,7 @@ def tasks(self, **kwargs):
         offset (deprecated) is the number of results to skip (for showing more pages).
         """
         allowed_kwargs = {'start_time', 'end_time', 'status', 'type', 'project',
-                          'batch', 'limit', 'offset', 'completed_before', 'completed_after', 
+                          'batch', 'limit', 'offset', 'completed_before', 'completed_after',
                           'next_token'}
         for key in kwargs:
             if key not in allowed_kwargs:
@@ -140,6 +151,29 @@ def create_task(self, task_type, **kwargs):
         taskdata = self._postrequest(endpoint, payload=kwargs)
         return Task(taskdata, self)
 
+    def create_batch(self, project, batch_name, callback):
+        payload = dict(project=project, name=batch_name, callback=callback)
+        batchdata = self._postrequest('batches', payload)
+        return Batch(batchdata, self)
+
+    def get_batch(self, batch_name: str):
+        batchdata = self._getrequest('batches/%s' % batch_name)
+        return Batch(batchdata, self)
+
+    def list_batches(self, **kwargs):
+        allowed_kwargs = { 'start_time', 'end_time', 'status', 'project',
+                           'batch', 'limit', 'offset', }
+        for key in kwargs:
+            if key not in allowed_kwargs:
+                raise ScaleInvalidRequest('Illegal parameter %s for ScaleClient.tasks()'
+                                          % key, None)
+        response = self._getrequest('tasks', params=kwargs)
+        docs = [Batch(doc, self) for doc in response['docs']]
+        return Batchlist(
+            docs, response['total'], response['limit'], response['offset'],
+            response['has_more'], response.get('next_token'),
+        )
+
 
 def _AddTaskTypeCreator(task_type):
     def create_task_wrapper(self, **kwargs):

scaleapi/batches.py

@@ -0,0 +1,26 @@
+class Batch(object):
+    def __init__(self, param_dict, client):
+        self.param_dict = param_dict
+        self.name = param_dict['name']
+        self.pending = None
+        self.completed = None
+        self.error = None
+        self.canceled = None
+        self.client = client
+
+    def __hash__(self):
+        return hash(self.name)
+
+    def __str__(self):
+        return 'Batch(name=%s)' % self.name
+
+    def __repr__(self):
+        return 'Batch(%s)' % self.param_dict
+
+    def finalize(self):
+        return self.client._postrequest("batches/%s/finalize" % self.name)
+
+    def get_status(self):
+        res = self.client._getrequest("batches/%s/status" % self.name)
+        for stat in ["pending", "completed", "error", "canceled"]:
+            setattr(self, stat, res.get(stat, 0))