Add peewee integration

fastapi_storages/integrations/peewee.py

@@ -0,0 +1,105 @@
+from typing import Any, Optional
+
+from peewee import CharField
+
+try:
+    from PIL import Image, UnidentifiedImageError
+
+    PIL = True
+except ImportError:  # pragma: no cover
+    PIL = False
+
+from fastapi_storages.base import BaseStorage, StorageFile, StorageImage
+from fastapi_storages.exceptions import ValidationException
+
+
+class FileType(CharField):
+    """
+    File type to be used with Storage classes. Stores the file name in the column.
+
+    ???+ usage
+        ```python
+        from fastapi_storages import FileSystemStorage
+        from fastapi_storages.integrations.peewee import FileType
+
+        class Example(Model):
+            file = FileType(storage=FileSystemStorage(path="/tmp"))
+        ```
+    """
+
+    def __init__(self, storage: BaseStorage, *args: Any, **kwargs: Any) -> None:
+        self.storage = storage
+        super().__init__(*args, **kwargs)
+
+    def db_value(self, value: Any) -> Optional[str]:
+        if value is None:
+            return value
+        if len(value.file.read(1)) != 1:
+            return None
+
+        file = StorageFile(name=value.filename, storage=self.storage)
+        file.write(file=value.file)
+
+        value.file.close()
+        return file.name
+
+    def python_value(self, value: Any) -> Optional[StorageFile]:
+        if value is None:
+            return value
+
+        return StorageFile(name=value, storage=self.storage)
+
+
+class ImageType(CharField):
+    """
+    Image type using `PIL` package to be used with Storage classes.
+    Stores the image path in the column.
+
+    ???+ usage
+        ```python
+        from fastapi_storages import FileSystemStorage
+        from fastapi_storages.integrations.peewee import ImageType
+
+        class Example(Model):
+            image = ImageType(storage=FileSystemStorage(path="/tmp"))
+        ```
+    """
+
+    def __init__(self, storage: BaseStorage, *args: Any, **kwargs: Any) -> None:
+        assert PIL is True, "'Pillow' package is required."
+
+        self.storage = storage
+        super().__init__(*args, **kwargs)
+
+    def db_value(self, value: Any) -> Optional[str]:
+        if value is None:
+            return value
+        if len(value.file.read(1)) != 1:
+            return None
+
+        try:
+            image_file = Image.open(value.file)
+            image_file.verify()
+        except UnidentifiedImageError:
+            raise ValidationException("Invalid image file")
+
+        image = StorageImage(
+            name=value.filename,
+            storage=self.storage,
+            height=image_file.height,
+            width=image_file.width,
+        )
+        image.write(file=value.file)
+
+        image_file.close()
+        value.file.close()
+        return image.name
+
+    def python_value(self, value: Any) -> Optional[StorageImage]:
+        if value is None:
+            return value
+
+        image = Image.open(self.storage.get_path(value))
+        return StorageImage(
+            name=value, storage=self.storage, height=image.height, width=image.width
+        )

pyproject.toml

@@ -36,6 +36,7 @@ dynamic = ["version"]
 full = [
   "Pillow~=9.4",
   "sqlalchemy>=1.4",
+  "peewee>=3",
 ]
 
 [project.urls]
@@ -64,6 +65,7 @@ dependencies = [
   "coverage==6.5.0",
   "moto==4.1.2",
   "mypy==0.982",
+  "peewee>=3",
   "Pillow==9.4.0",
   "pytest==7.2.0",
   "ruff==0.0.237",

tests/test_integrations/test_peewee/test_file.py

@@ -0,0 +1,57 @@
+import io
+from pathlib import Path
+
+import pytest
+from peewee import AutoField, Model, SqliteDatabase
+
+from fastapi_storages import FileSystemStorage
+from fastapi_storages.integrations.peewee import FileType
+from tests.engine import database_uri
+from tests.test_integrations.utils import UploadFile
+
+db = SqliteDatabase(database_uri)
+
+
+class Model(Model):
+    id = AutoField(primary_key=True)
+    file = FileType(storage=FileSystemStorage(path="/tmp"), null=True)
+
+    class Meta:
+        database = db
+
+
+@pytest.fixture(autouse=True)
+def prepare_database():
+    db.create_tables([Model])
+    yield
+    db.drop_tables([Model])
+
+
+def test_valid_file(tmp_path: Path) -> None:
+    Model.file.storage = FileSystemStorage(path=str(tmp_path))
+
+    input_file = tmp_path / "input.txt"
+    input_file.write_bytes(b"123")
+
+    upload_file = UploadFile(file=input_file.open("rb"), filename="example.txt")
+    Model.create(file=upload_file)
+    model = Model.get()
+
+    assert model.file.name == "example.txt"
+    assert model.file.size == 3
+    assert model.file.path == str(tmp_path / "example.txt")
+
+
+def test_nullable_file() -> None:
+    model = Model(file=None)
+    model.save()
+
+    assert model.file is None
+
+
+def test_clear_empty_file() -> None:
+    upload_file = UploadFile(file=io.BytesIO(b""), filename="")
+    Model.create(file=upload_file)
+
+    model = Model.get()
+    assert model.file is None

tests/test_integrations/test_peewee/test_image.py

@@ -0,0 +1,69 @@
+import io
+from pathlib import Path
+
+import pytest
+from peewee import AutoField, Model, SqliteDatabase
+from PIL import Image
+
+from fastapi_storages import FileSystemStorage
+from fastapi_storages.exceptions import ValidationException
+from fastapi_storages.integrations.peewee import ImageType
+from tests.engine import database_uri
+from tests.test_integrations.utils import UploadFile
+
+db = SqliteDatabase(database_uri)
+
+
+class Model(Model):
+    id = AutoField(primary_key=True)
+    image = ImageType(storage=FileSystemStorage(path="/tmp"), null=True)
+
+    class Meta:
+        database = db
+
+
+@pytest.fixture(autouse=True)
+def prepare_database():
+    db.create_tables([Model])
+    yield
+    db.drop_tables([Model])
+
+
+def test_valid_image(tmp_path: Path) -> None:
+    Model.image.storage = FileSystemStorage(path=str(tmp_path))
+
+    input_file = tmp_path / "input.png"
+    image = Image.new("RGB", (800, 1280), (255, 255, 255))
+    image.save(input_file, "PNG")
+
+    upload_file = UploadFile(file=input_file.open("rb"), filename="image.png")
+    Model.create(image=upload_file)
+    model = Model.get()
+
+    assert model.image.name == "image.png"
+    assert model.image.size == 5847
+    assert model.image.path == str(tmp_path / "image.png")
+
+
+def test_invalid_image(tmp_path: Path) -> None:
+    input_file = tmp_path / "image.png"
+    input_file.write_bytes(b"123")
+    upload_file = UploadFile(file=input_file.open("rb"), filename="image.png")
+
+    with pytest.raises(ValidationException):
+        Model.create(image=upload_file)
+
+
+def test_nullable_image() -> None:
+    Model.create(image=None)
+    model = Model.get()
+
+    assert model.image is None
+
+
+def test_clear_empty_image() -> None:
+    upload_file = UploadFile(file=io.BytesIO(b""), filename="")
+    Model.create(image=upload_file)
+    model = Model.get()
+
+    assert model.image is None

tests/test_integrations/test_sqlalchemy/test_file.py

@@ -1,6 +1,5 @@
 import io
 from pathlib import Path
-from typing import BinaryIO
 
 import pytest
 from sqlalchemy import Column, Integer, create_engine
@@ -9,21 +8,12 @@
 from fastapi_storages import FileSystemStorage
 from fastapi_storages.integrations.sqlalchemy import FileType
 from tests.engine import database_uri
+from tests.test_integrations.utils import UploadFile
 
 Base = declarative_base()
 engine = create_engine(database_uri)
 
 
-class UploadFile:
-    """
-    Dummy UploadFile like the one in Starlette.
-    """
-
-    def __init__(self, file: BinaryIO, filename: str) -> None:
-        self.file = file
-        self.filename = filename
-
-
 class Model(Base):
     __tablename__ = "model"
 

tests/test_integrations/test_sqlalchemy/test_image.py

@@ -1,6 +1,5 @@
 import io
 from pathlib import Path
-from typing import BinaryIO
 
 import pytest
 from PIL import Image
@@ -11,21 +10,12 @@
 from fastapi_storages import FileSystemStorage
 from fastapi_storages.integrations.sqlalchemy import ImageType
 from tests.engine import database_uri
+from tests.test_integrations.utils import UploadFile
 
 Base = declarative_base()
 engine = create_engine(database_uri)
 
 
-class UploadFile:
-    """
-    Dummy UploadFile like the one in Starlette.
-    """
-
-    def __init__(self, file: BinaryIO, filename: str) -> None:
-        self.file = file
-        self.filename = filename
-
-
 class Model(Base):
     __tablename__ = "model"
 

tests/test_integrations/utils.py

@@ -0,0 +1,11 @@
+from typing import BinaryIO
+
+
+class UploadFile:
+    """
+    Dummy UploadFile like the one in Starlette.
+    """
+
+    def __init__(self, file: BinaryIO, filename: str) -> None:
+        self.file = file
+        self.filename = filename