diff --git a/README.rst b/README.rst index d4a2fa9..2f5403e 100644 --- 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 index 79ff89f..23e4a21 100644 --- 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 index 5a76857..1319c8a 100644 --- 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 new file mode 100644 index 0000000..1ebca22 --- /dev/null +++ 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 index d956587..a811469 100644 --- 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 new file mode 100644 index 0000000..ac0a33c Binary files /dev/null and b/tests/test_image.png differ