diff --git a/src/access.py b/src/access.py index 322d28b..896c115 100644 --- a/src/access.py +++ b/src/access.py @@ -3,6 +3,7 @@ import os import time import requests from threading import Event +from io import BytesIO class Access: @@ -111,39 +112,18 @@ class Access: return content - def send_image(self, img, doc_id): - if self.is_time_for_refreshing(): - self.refresh_token() - - self.refreshing.wait() - - with open(img, 'rb') as img_file: - files = { - 'image': (os.path.basename(img), img_file) - } - response = requests.post(f'{self.url}/doc-convert/image', files=files, headers=self.headers) - - if response.status_code == 400: - self.update_status(doc_id, self.ERROR) - raise Exception(f'400 Bad request: {response.json()["message"]}.') - elif response.status_code == 200: - img_url = response.json()['imageUrl'] - else: - self.update_status(doc_id, self.ERROR) - raise Exception(f'{response.status_code}') - - return img_url - - def send_image_by_bytes(self, img_file_path, img_content: bytes, doc_id): + def send_image(self, img_path, doc_id, img_content: bytes = None): if self.is_time_for_refreshing(): self.refresh_token() self.refreshing.wait() + img_obj = BytesIO(img_content) if img_content else open(img_path, 'rb') files = { - 'image': (os.path.basename(img_file_path), img_content) + 'image': (os.path.basename(img_path), img_obj) } response = requests.post(f'{self.url}/doc-convert/image', files=files, headers=self.headers) + img_obj.close() if response.status_code == 400: self.update_status(doc_id, self.ERROR) diff --git a/src/html_epub_preprocessor.py b/src/html_epub_preprocessor.py index d486d43..baf1b59 100644 --- a/src/html_epub_preprocessor.py +++ b/src/html_epub_preprocessor.py @@ -23,7 +23,7 @@ def save_image_locally(img_file_path, img_content, book_id): def save_image_to_aws(access: Access, img_file_path, img_content: bytes, book_id): - link = access.send_image_by_bytes(img_file_path, img_content, book_id) + link = access.send_image(img_file_path, doc_id=book_id, img_content=img_content) return link diff --git a/src/html_preprocessor.py b/src/html_preprocessor.py index c82607d..b01298c 100644 --- a/src/html_preprocessor.py +++ b/src/html_preprocessor.py @@ -380,7 +380,7 @@ class HTMLPreprocessor: img_path = pathlib.Path(f'{html_path.parent}', f'{img_name}') if access is not None: - link = access.send_image(img_path, book_id) + link = access.send_image(img_path, doc_id=book_id) img.attrs['src'] = link self.logger_object.log(f'{img_name} successfully uploaded.') else: