Add support for files endpoint (#38)

* Add support for files endpoint * add readme * fix readme * nit * simplify upload endpoint * Extended Readme with a task payload example Co-authored-by: Fatih Kurtoglu <[email protected]>
scaleapi · May 17, 2021 · 727bb95 · 727bb95
1 parent 1ce41cd
commit 727bb95
Show file tree

Hide file tree

Showing 6 changed files with 178 additions and 8 deletions.
diff --git a/README.rst b/README.rst
@@ -337,6 +337,60 @@ __ https://docs.scale.com/reference#project-update-parameters
         instruction="update: Please label all the stuff",
     )
 
+Files
+________
+
+Files are a way of uploading local files directly to Scale storage or importing files before creating tasks.
+
+The ``file.attachment_url`` can be used in place of attachments in task payload.
+
+Upload Files
+^^^^^^^^^^^^^^
+
+Upload a file. Check out `Scale's API documentation`__ for more information.
+
+__ https://docs.scale.com/reference#file-upload-1
+
+.. code-block:: python
+
+    with open(file_name, 'rb') as f:
+        my_file = client.upload_file(
+            file=f,
+            project_name = "test_project",
+        )
+
+Import Files
+^^^^^^^^^^^^^^
+
+Import a file from a URL. Check out `Scale's API documentation`__ for more information.
+
+__ https://docs.scale.com/reference#file-import-1
+
+.. code-block:: python
+
+    my_file = client.import_file(
+        file_url="http://i.imgur.com/v4cBreD.jpg",
+        project_name = "test_project",
+    )
+
+
+After the files are successfully uploaded to Scale's storage, you can access the URL as ``my_file.attachment_url``, which will have a prefix like ``scaledata://``.
+
+The attribute can be passed to the task payloads, in the ``attachment`` parameter.
+
+.. code-block:: python
+
+  task_payload = dict(
+      ...
+      ...
+      attachment_type = "image",
+      attachment = my_file.attachment_url,
+      ...
+      ...
+  )
+
+
+
 Error handling
 ______________
 

diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py
@@ -1,7 +1,8 @@
-from typing import Dict, Generator, Generic, List, TypeVar, Union
+from typing import IO, Dict, Generator, Generic, List, TypeVar, Union
 
 from scaleapi.batches import Batch, BatchStatus
 from scaleapi.exceptions import ScaleInvalidRequest
+from scaleapi.files import File
 from scaleapi.projects import Project
 
 from ._version import __version__  # noqa: F401
@@ -317,7 +318,13 @@ def create_task(self, task_type: TaskType, **kwargs) -> Task:
         taskdata = self.api.post_request(endpoint, body=kwargs)
         return Task(taskdata, self)
 
-    def create_batch(self, project: str, batch_name: str, callback: str = "") -> Batch:
+    def create_batch(
+        self,
+        project: str,
+        batch_name: str,
+        callback: str = "",
+        instruction_batch: bool = False,
+    ) -> Batch:
         """Create a new Batch within a project.
         https://docs.scale.com/reference#batch-creation
 
@@ -329,12 +336,21 @@ def create_batch(self, project: str, batch_name: str, callback: str = "") -> Bat
             callback (str, optional):
                 Email to notify, or URL to POST to
                 when a batch is complete.
+            instruction_batch (bool):
+                Only applicable for self serve projects.
+                Create an instruction batch by setting
+                the instruction_batch flag to true.
 
         Returns:
             Batch: Created batch object
         """
         endpoint = "batches"
-        payload = dict(project=project, name=batch_name, callback=callback)
+        payload = dict(
+            project=project,
+            name=batch_name,
+            instruction_batch=instruction_batch,
+            callback=callback,
+        )
         batchdata = self.api.post_request(endpoint, body=payload)
         return Batch(batchdata, self)
 
@@ -596,3 +612,39 @@ def update_project(self, project_name: str, **kwargs) -> Project:
         endpoint = f"projects/{Api.quote_string(project_name)}/setParams"
         projectdata = self.api.post_request(endpoint, body=kwargs)
         return Project(projectdata, self)
+
+    def upload_file(self, file: IO, **kwargs) -> File:
+        """Upload file.
+        Refer to Files API Reference:
+        https://docs.scale.com/reference#file-upload-1
+
+        Args:
+            file (IO):
+                File buffer
+
+        Returns:
+            File
+        """
+
+        endpoint = "files/upload"
+        files = {"file": file}
+        filedata = self.api.post_request(endpoint, files=files, data=kwargs)
+        return File(filedata, self)
+
+    def import_file(self, file_url: str, **kwargs) -> File:
+        """Import file from a remote url.
+        Refer to Files API Reference:
+        https://docs.scale.com/reference#file-import-1
+
+        Args:
+            file_url (str):
+                File's url
+
+        Returns:
+            File
+        """
+
+        endpoint = "files/import"
+        payload = dict(file_url=file_url, **kwargs)
+        filedata = self.api.post_request(endpoint, body=payload)
+        return File(filedata, self)
diff --git a/scaleapi/api.py b/scaleapi/api.py
@@ -30,10 +30,20 @@ def __init__(self, api_key, user_agent_extension=None):
             "Content-Type": "application/json",
             "User-Agent": self._generate_useragent(user_agent_extension),
         }
+        self._headers_multipart_form_data = {
+            "User-Agent": self._generate_useragent(user_agent_extension),
+        }
 
     @staticmethod
     def _http_request(
-        method, url, headers=None, auth=None, params=None, body=None
+        method,
+        url,
+        headers=None,
+        auth=None,
+        params=None,
+        body=None,
+        files=None,
+        data=None,
     ) -> Response:
 
         https = requests.Session()
@@ -59,6 +69,8 @@ def _http_request(
                 auth=auth,
                 params=params,
                 json=body,
+                files=files,
+                data=data,
             )
 
             return res
@@ -77,13 +89,21 @@ def _raise_on_respose(res: Response):
         raise exception(message, res.status_code)
 
     def _api_request(
-        self, method, endpoint, headers=None, auth=None, params=None, body=None
+        self,
+        method,
+        endpoint,
+        headers=None,
+        auth=None,
+        params=None,
+        body=None,
+        files=None,
+        data=None,
     ):
         """Generic HTTP request method with error handling."""
 
         url = f"{SCALE_ENDPOINT}/{endpoint}"
 
-        res = self._http_request(method, url, headers, auth, params, body)
+        res = self._http_request(method, url, headers, auth, params, body, files, data)
 
         json = None
         if res.status_code == 200:
@@ -99,10 +119,18 @@ def get_request(self, endpoint, params=None):
             "GET", endpoint, headers=self._headers, auth=self._auth, params=params
         )
 
-    def post_request(self, endpoint, body=None):
+    def post_request(self, endpoint, body=None, files=None, data=None):
         """Generic POST Request Wrapper"""
         return self._api_request(
-            "POST", endpoint, headers=self._headers, auth=self._auth, body=body
+            "POST",
+            endpoint,
+            headers=self._headers
+            if files is None
+            else self._headers_multipart_form_data,
+            auth=self._auth,
+            body=body,
+            files=files,
+            data=data,
         )
 
     @staticmethod

diff --git a/scaleapi/files.py b/scaleapi/files.py
@@ -0,0 +1,21 @@
+class File:
+    """File class, containing File information."""
+
+    def __init__(self, json, client):
+        self._json = json
+        self.id = json["id"]
+        self.attachment_url = json["attachment_url"]
+        self._client = client
+
+    def __hash__(self):
+        return hash(self.id)
+
+    def __str__(self):
+        return f"File(id={self.id})"
+
+    def __repr__(self):
+        return f"File({self._json})"
+
+    def as_dict(self):
+        """Returns all attributes as a dictionary"""
+        return self._json
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -378,3 +378,18 @@ def test_get_batches():
     # Download all batches to check total count
     all_batches = list(client.get_batches(project_name=TEST_PROJECT_NAME))
     assert total_batches == len(all_batches)
+
+
+def test_files_upload():
+    with open("tests/test_image.png", "rb") as f:
+        client.upload_file(
+            file=f,
+            project_name=TEST_PROJECT_NAME,
+        )
+
+
+def test_files_import():
+    client.import_file(
+        file_url="https://static.scale.com/uploads/selfserve-sample-image.png",
+        project_name=TEST_PROJECT_NAME,
+    )
diff --git a/tests/test_image.png b/tests/test_image.png