using PIL for determining media type

Author: Eldies

app/storage_manager.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
-import json
 import os
+from dataclasses import dataclass
+
+from PIL import Image as PILImage
 
 from .settings import settings
 
@@ -9,37 +11,37 @@ class StorageManagerException(Exception):
     pass
 
 
+@dataclass
+class Image:
+    data: bytes
+    mimetype: str
+
+
 class StorageManager(object):
     def path_for_uuid(self, uuid: list[str]) -> str:
         return os.path.join(settings.upload_folder, *uuid)
 
     def uuid_exists(self, uuid: list[str]) -> bool:
         return os.path.exists(self.path_for_uuid(uuid))
 
-    def save_image(self, uuid: list[str], file_content: bytes, data: str | None = None) -> None:
+    def save_image(self, uuid: list[str], file_content: bytes) -> None:
         folder = self.path_for_uuid(uuid)
         os.makedirs(folder)
         with open(os.path.join(folder, "file"), "wb") as f:
             f.write(file_content)
-        if data:
-            with open(os.path.join(folder, "data"), "w") as f:
-                f.write(data)
-
-    def _read_data(self, uuid: list[str]) -> dict[str, str]:
-        path = os.path.join(self.path_for_uuid(uuid), "data")
-        if not os.path.exists(path):
-            return {}
-        with open(path, "r") as f:
-            return json.loads(f.read())
-
-    def _read_file(self, uuid: list[str]) -> bytes:
-        with open(os.path.join(self.path_for_uuid(uuid), "file"), "rb") as f:
-            return f.read()
-
-    def get_file(self, uuid: list[str]) -> tuple[bytes, dict]:
+
+    def get_image(self, uuid: list[str]) -> Image:
         if not self.uuid_exists(uuid):
             raise StorageManagerException("Not Found")
-        return self._read_file(uuid), self._read_data(uuid)
+        filename = os.path.join(self.path_for_uuid(uuid), "file")
+        with open(filename, "rb") as f:
+            data = f.read()
+        pil_image = PILImage.open(filename)
+        assert pil_image.format is not None
+        return Image(
+            data=data,
+            mimetype=PILImage.MIME[pil_image.format],
+        )
 
 
 def get_storage_manager() -> StorageManager:

app/views.py

@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 import base64
-import json
 import logging
 
 from fastapi import APIRouter, HTTPException
@@ -50,11 +49,6 @@ async def post_image(
     get_storage_manager().save_image(
         uuid=[client.id, filename],
         file_content=base64.b64decode(params.base64),
-        data=json.dumps(
-            dict(
-                mimetype="image/jpeg",
-            )
-        ),
     )
 
     return dict(
@@ -66,8 +60,8 @@ async def post_image(
 @router_api.get("/image/{client_id}/{uuid}")
 async def get_image(client_id: str, uuid: str) -> Response:
     try:
-        content, data = get_storage_manager().get_file([client_id, uuid])
-        return Response(content=content, media_type=data.get("mimetype"))
+        image = get_storage_manager().get_image([client_id, uuid])
+        return Response(content=image.data, media_type=image.mimetype)
     except StorageManagerException as e:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,

pyproject.toml

@@ -10,6 +10,7 @@ fastapi = "~0.111"
 uvicorn = "~0.30"
 base58 = "2.1.1"
 pydantic-settings = "~2.2.1"
+pillow = "~10.3.0"
 
 [tool.poetry.group.test]
 optional = true

tests/test_storage_manager.py

@@ -1,7 +1,9 @@
 # -*- coding: utf-8 -*-
+import io
 import os
 
 import pytest
+from PIL import Image
 
 from app import storage_manager
 from app.settings import Settings
@@ -14,26 +16,22 @@ def _setup(self, upload_folder, monkeypatch):
         self.upload_folder = upload_folder
         self.manager = storage_manager.StorageManager()
 
-    def check_save(self, uuid, data=None):
+        self.image = Image.new(mode="RGB", size=(3, 3))
+        img_stream = io.BytesIO()
+        self.image.save(img_stream, format="jpeg")
+        self.image_byte_array = img_stream.getvalue()
+
+    def check_save(self, uuid):
         folder = os.path.join(self.upload_folder, uuid)
         file_path = os.path.join(folder, "file")
-        data_path = os.path.join(folder, "data")
 
         assert os.path.isdir(folder)
         assert os.path.isfile(file_path)
         with open(file_path, "rb") as file:
-            assert file.read() == b"abcdef"
-        if data is not None:
-            assert os.path.isfile(data_path)
-            with open(data_path, "r") as file:
-                assert file.read() == data
+            assert file.read() == self.image_byte_array
 
     def test_save_image_ok(self):
-        self.manager.save_image(["some_uuid"], b"abcdef", "some_data")
-        self.check_save("some_uuid", "some_data")
-
-    def test_save_image_no_data_ok(self):
-        self.manager.save_image(["some_uuid"], b"abcdef")
+        self.manager.save_image(["some_uuid"], self.image_byte_array)
         self.check_save("some_uuid")
 
     @pytest.mark.parametrize("exists", [True, False])
@@ -42,17 +40,13 @@ def test_exists(self, exists):
             os.makedirs(os.path.join(self.upload_folder, "some_uuid"))
         assert self.manager.uuid_exists(["some_uuid"]) == exists
 
-    def test_read_data(self):
-        os.makedirs(os.path.join(self.upload_folder, "some_uuid"))
-        with open(os.path.join(self.upload_folder, "some_uuid", "data"), "w") as f:
-            f.write('{"foo": "bar"}')
-        assert self.manager._read_data(["some_uuid"]) == {"foo": "bar"}
-
-    def test_read_file(self):
+    def test_get_image(self):
         os.makedirs(os.path.join(self.upload_folder, "some_uuid"))
         with open(os.path.join(self.upload_folder, "some_uuid", "file"), "wb") as f:
-            f.write(b"READ DATA")
-        assert self.manager._read_file(["some_uuid"]) == b"READ DATA"
+            f.write(self.image_byte_array)
+        result_image = self.manager.get_image(["some_uuid"])
+        assert result_image.data == self.image_byte_array
+        assert result_image.mimetype == "image/jpeg"
 
     def test_path_for_uuid_with_sep(self):
         path = self.manager.path_for_uuid(["foo", "bar"])

tests/views/test_images.py

@@ -1,10 +1,11 @@
 # -*- coding: utf-8 -*-
 import base64
-import json
+import io
 import os
 
 import pytest
 import pytest_asyncio
+from PIL import Image
 
 
 @pytest.fixture()
@@ -63,28 +64,28 @@ async def test_ok(self, filename):
         assert os.path.isdir(os.path.join(self.upload_folder, "test_client", expected_filename))
         with open(os.path.join(self.upload_folder, "test_client", expected_filename, "file"), "rb") as file:
             assert file.read() == b"abcdef"
-        with open(os.path.join(self.upload_folder, "test_client", expected_filename, "data"), "r") as file:
-            assert json.loads(file.read()) == dict(
-                mimetype="image/jpeg",
-            )
 
 
 @pytest.mark.asyncio
 class TestImageViewGet:
     @pytest_asyncio.fixture(autouse=True)
     def _setup(self, client, upload_folder):
         self.client = client
+
+        image = Image.new(mode="RGB", size=(3, 3))
+        img_stream = io.BytesIO()
+        image.save(img_stream, format="jpeg")
+        self.image_byte_array = img_stream.getvalue()
+
         os.makedirs(os.path.join(upload_folder, "some_client_id", "some_uuid"))
         with open(os.path.join(upload_folder, "some_client_id", "some_uuid", "file"), "wb") as f:
-            f.write(b"abcdef")
-        with open(os.path.join(upload_folder, "some_client_id", "some_uuid", "data"), "w") as f:
-            f.write(json.dumps(dict(mimetype="image/jpeg")))
+            f.write(self.image_byte_array)
 
     async def test_ok(self):
         response = await self.client.get("/v1/image/some_client_id/some_uuid")
         assert response.status_code == 200
         assert response.headers["Content-Type"] == "image/jpeg"
-        assert response.content == b"abcdef"
+        assert response.content == self.image_byte_array
 
     async def test_no_client_id(self):
         response = await self.client.get("/v1/image/some_uuid")