fixes for #199 and unreported choices bug

docs/releases.md

@@ -3,6 +3,9 @@
 ## 🐛 Fixes
 
 * Fix populating default values in pk_only child models [#202](https://github.com/collerek/ormar/issues/202)
+* Fix mypy for LargeBinary fields with base64 str representation [#199](https://github.com/collerek/ormar/issues/199)
+* Fix OpenAPI schema format for LargeBinary fields with base64 str representation [#199](https://github.com/collerek/ormar/issues/199)
+* Fix OpenAPI choices encoding for LargeBinary fields with base64 str representation
 
 # 0.10.7
 

ormar/fields/model_fields.py

@@ -1,7 +1,7 @@
 import datetime
 import decimal
 import uuid
-from typing import Any, Optional, TYPE_CHECKING
+from typing import Any, List, Literal, Optional, TYPE_CHECKING, Union, overload
 
 import pydantic
 import sqlalchemy
@@ -426,52 +426,84 @@ class JSON(ModelFieldFactory, pydantic.Json):
         return sqlalchemy.JSON()
 
 
-class LargeBinary(ModelFieldFactory, bytes):
+if TYPE_CHECKING:  # pragma: nocover
-    """
-    LargeBinary field factory that construct Field classes and populated their values.
-    """
 
-    _type = bytes
+    @overload
-    _sample = "bytes"
+    def LargeBinary(
+        max_length: int, *, represent_as_base64_str: Literal[True], **kwargs
+    ) -> str:
+        ...
 
-    def __new__(  # type: ignore # noqa CFQ002
+    @overload
-        cls, *, max_length: int, represent_as_base64_str: bool = False, **kwargs: Any
+    def LargeBinary(
-    ) -> BaseField:  # type: ignore
+        max_length: int, *, represent_as_base64_str: Literal[False], **kwargs
-        kwargs = {
+    ) -> bytes:
-            **kwargs,
+        ...
-            **{
-                k: v
-                for k, v in locals().items()
-                if k not in ["cls", "__class__", "kwargs"]
-            },
-        }
-        return super().__new__(cls, **kwargs)
 
-    @classmethod
+    @overload
-    def get_column_type(cls, **kwargs: Any) -> Any:
+    def LargeBinary(
+        max_length: int, represent_as_base64_str: Literal[False] = ..., **kwargs
+    ) -> bytes:
+        ...
+
+    def LargeBinary(
+        max_length: int, represent_as_base64_str: bool = False, **kwargs: Any
+    ) -> Union[str, bytes]:
+        pass
+
+
+else:
+
+    class LargeBinary(ModelFieldFactory, bytes):
+        """
+        LargeBinary field factory that construct Field classes and populated their values.
         """
-        Return proper type of db column for given field type.
-        Accepts required and optional parameters that each column type accepts.
 
-        :param kwargs: key, value pairs of sqlalchemy options
+        _type = bytes
-        :type kwargs: Any
+        _sample = "bytes"
-        :return: initialized column with proper options
-        :rtype: sqlalchemy Column
-        """
-        return sqlalchemy.LargeBinary(length=kwargs.get("max_length"))
 
-    @classmethod
+        def __new__(  # type: ignore # noqa CFQ002
-    def validate(cls, **kwargs: Any) -> None:
+            cls,
-        """
+            *,
-        Used to validate if all required parameters on a given field type are set.
+            max_length: int,
-        :param kwargs: all params passed during construction
+            represent_as_base64_str: bool = False,
-        :type kwargs: Any
+            **kwargs: Any
-        """
+        ) -> BaseField:  # type: ignore
-        max_length = kwargs.get("max_length", None)
+            kwargs = {
-        if max_length <= 0:
+                **kwargs,
-            raise ModelDefinitionError(
+                **{
-                "Parameter max_length is required for field LargeBinary"
+                    k: v
-            )
+                    for k, v in locals().items()
+                    if k not in ["cls", "__class__", "kwargs"]
+                },
+            }
+            return super().__new__(cls, **kwargs)
+
+        @classmethod
+        def get_column_type(cls, **kwargs: Any) -> Any:
+            """
+            Return proper type of db column for given field type.
+            Accepts required and optional parameters that each column type accepts.
+
+            :param kwargs: key, value pairs of sqlalchemy options
+            :type kwargs: Any
+            :return: initialized column with proper options
+            :rtype: sqlalchemy Column
+            """
+            return sqlalchemy.LargeBinary(length=kwargs.get("max_length"))
+
+        @classmethod
+        def validate(cls, **kwargs: Any) -> None:
+            """
+            Used to validate if all required parameters on a given field type are set.
+            :param kwargs: all params passed during construction
+            :type kwargs: Any
+            """
+            max_length = kwargs.get("max_length", None)
+            if max_length <= 0:
+                raise ModelDefinitionError(
+                    "Parameter max_length is required for field LargeBinary"
+                )
 
 
 class BigInteger(Integer, int):

ormar/models/helpers/validation.py

@@ -142,7 +142,10 @@ def generate_model_example(model: Type["Model"], relation_map: Dict = None) -> D
     )
     for name, field in model.Meta.model_fields.items():
         if not field.is_relation:
-            example[name] = field.__sample__
+            if field.__type__ == bytes and field.represent_as_base64_str:
+                example[name] = "string"
+            else:
+                example[name] = field.__sample__
         elif isinstance(relation_map, dict) and name in relation_map:
             example[name] = get_nested_model_example(
                 name=name, field=field, relation_map=relation_map
@@ -217,6 +220,44 @@ def get_pydantic_example_repr(type_: Any) -> Any:
         return "string"
 
 
+def overwrite_example_and_description(
+    schema: Dict[str, Any], model: Type["Model"]
+) -> None:
+    """
+    Overwrites the example with properly nested children models.
+    Overwrites the description if it's taken from ormar.Model.
+
+    :param schema: schema of current model
+    :type schema: Dict[str, Any]
+    :param model: model class
+    :type model: Type["Model"]
+    """
+    schema["example"] = generate_model_example(model=model)
+    if "Main base class of ormar Model." in schema.get("description", ""):
+        schema["description"] = f"{model.__name__}"
+
+
+def overwrite_binary_format(schema: Dict[str, Any], model: Type["Model"]) -> None:
+    """
+    Overwrites format of the field if it's a LargeBinary field with
+    a flag to represent the field as base64 encoded string.
+
+    :param schema: schema of current model
+    :type schema: Dict[str, Any]
+    :param model: model class
+    :type model: Type["Model"]
+    """
+    for field_id, prop in schema.get("properties", {}).items():
+        if (
+            field_id in model._bytes_fields
+            and model.Meta.model_fields[field_id].represent_as_base64_str
+        ):
+            prop["format"] = "base64"
+            prop["enum"] = [
+                base64.b64encode(choice).decode() for choice in prop["enum"]
+            ]
+
+
 def construct_modify_schema_function(fields_with_choices: List) -> SchemaExtraCallable:
     """
     Modifies the schema to include fields with choices validator.
@@ -237,9 +278,8 @@ def construct_modify_schema_function(fields_with_choices: List) -> SchemaExtraCa
             if field_id in fields_with_choices:
                 prop["enum"] = list(model.Meta.model_fields[field_id].choices)
                 prop["description"] = prop.get("description", "") + "An enumeration."
-        schema["example"] = generate_model_example(model=model)
+        overwrite_example_and_description(schema=schema, model=model)
-        if "Main base class of ormar Model." in schema.get("description", ""):
+        overwrite_binary_format(schema=schema, model=model)
-            schema["description"] = f"{model.__name__}"
 
     return staticmethod(schema_extra)  # type: ignore
 
@@ -256,9 +296,8 @@ def construct_schema_function_without_choices() -> SchemaExtraCallable:
     """
 
     def schema_extra(schema: Dict[str, Any], model: Type["Model"]) -> None:
-        schema["example"] = generate_model_example(model=model)
+        overwrite_example_and_description(schema=schema, model=model)
-        if "Main base class of ormar Model." in schema.get("description", ""):
+        overwrite_binary_format(schema=schema, model=model)
-            schema["description"] = f"{model.__name__}"
 
     return staticmethod(schema_extra)  # type: ignore
 

tests/test_fastapi/test_binary_fields.py

@@ -52,7 +52,7 @@ class BinaryThing(ormar.Model):
 
     id: uuid.UUID = ormar.UUID(primary_key=True, default=uuid.uuid4)
     name: str = ormar.Text(default="")
-    bt: bytes = ormar.LargeBinary(
+    bt: str = ormar.LargeBinary(
         max_length=1000,
         choices=[blob3, blob4, blob5, blob6],
         represent_as_base64_str=True,
@@ -89,3 +89,14 @@ def test_read_main():
         assert response.json()[0]["bt"] == base64.b64encode(blob3).decode()
         thing = BinaryThing(**response.json()[0])
         assert thing.__dict__["bt"] == blob3
+
+
+def test_schema():
+    schema = BinaryThing.schema()
+    assert schema["properties"]["bt"]["format"] == "base64"
+    converted_choices = ["7g==", "/w==", "8CiMKA==", "wyg="]
+    assert len(schema["properties"]["bt"]["enum"]) == 4
+    assert all(
+        choice in schema["properties"]["bt"]["enum"] for choice in converted_choices
+    )
+    assert schema["example"]["bt"] == "string"