Fix byte to binary type mapping (#3114)

* Fix byte to binary type mapping

* Fix type mapping guard

Author: koxudaxi

src/datamodel_code_generator/model/pydantic_v2/imports.py

@@ -16,6 +16,7 @@
 IMPORT_FUTURE_DATETIME = Import.from_full_path("pydantic.FutureDatetime")
 IMPORT_PAST_DATE = Import.from_full_path("pydantic.PastDate")
 IMPORT_FUTURE_DATE = Import.from_full_path("pydantic.FutureDate")
+IMPORT_BASE64BYTES = Import.from_full_path("pydantic.Base64Bytes")
 IMPORT_BASE64STR = Import.from_full_path("pydantic.Base64Str")
 IMPORT_SERIALIZE_AS_ANY = Import.from_full_path("pydantic.SerializeAsAny")
 IMPORT_PYDANTIC_DATACLASS = Import.from_full_path("pydantic.dataclasses.dataclass")

src/datamodel_code_generator/model/pydantic_v2/types.py

@@ -27,6 +27,7 @@
 from datamodel_code_generator.model.pydantic_v2.imports import (
     IMPORT_ANYURL,
     IMPORT_AWARE_DATETIME,
+    IMPORT_BASE64BYTES,
     IMPORT_BASE64STR,
     IMPORT_CONBYTES,
     IMPORT_CONDECIMAL,
@@ -365,8 +366,11 @@ def get_data_str_type(self, types: Types, **kwargs: Any) -> DataType:
 
     def get_data_bytes_type(self, types: Types, **kwargs: Any) -> DataType:
         """Get bytes data type with constraints (conbytes)."""
+        base64_encoded = kwargs.pop("base64_encoded", False)
         data_type_kwargs: dict[str, Any] = self.transform_kwargs(kwargs, bytes_kwargs)
         strict = StrictTypes.bytes in self.strict_types
+        if base64_encoded and not data_type_kwargs and not strict:
+            return self.data_type.from_import(IMPORT_BASE64BYTES)
         if data_type_kwargs and not strict:
             return self.data_type.from_import(IMPORT_CONBYTES, kwargs=data_type_kwargs)
         # conbytes doesn't accept strict argument

src/datamodel_code_generator/parser/jsonschema.py

@@ -746,6 +746,11 @@ def _get_type_with_mappings(self, type_: str, format_: str | None = None) -> Typ
 
         return _get_type(type_, format_, data_formats)
 
+    def _is_base64_encoded_binary_mapping(self, type_: str, format_: str) -> bool:
+        if type_ != "string" or format_ != "byte" or not self.type_mappings:
+            return False
+        return self.type_mappings.get((type_, format_)) == "binary"
+
     @cached_property
     def schema_paths(self) -> list[tuple[str, list[str]]]:
         """Get schema paths for definitions and defs.
@@ -1292,8 +1297,12 @@ def _get_data_type(type_: str, format__: str) -> DataType:
             else:
                 kwargs_to_pass = obj.model_dump()
 
+            types = self._get_type_with_mappings(type_, format__)
+            if types == Types.binary and self._is_base64_encoded_binary_mapping(type_, format__):
+                kwargs_to_pass["base64_encoded"] = True
+
             return self.data_type_manager.get_data_type(
-                self._get_type_with_mappings(type_, format__),
+                types,
                 field_constraints=self.field_constraints,
                 **kwargs_to_pass,
             )

tests/data/expected/main/openapi/type_mappings_byte_to_binary.py

@@ -0,0 +1,11 @@
+# generated by datamodel-codegen:
+#   filename:  type_mappings_byte_to_binary.yaml
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import Base64Bytes, BaseModel
+
+
+class Foo(BaseModel):
+    bar: Base64Bytes

tests/data/openapi/type_mappings_byte_to_binary.yaml

@@ -0,0 +1,11 @@
+openapi: 3.0.3
+components:
+  schemas:
+    Foo:
+      type: object
+      properties:
+        bar:
+          type: string
+          format: byte
+      required:
+        - bar

tests/main/openapi/test_main_openapi.py

@@ -3855,6 +3855,23 @@ def test_main_openapi_byte_format(output_file: Path) -> None:
     )
 
 
+def test_main_openapi_type_mappings_byte_to_binary(output_file: Path) -> None:
+    """Test mapping OpenAPI byte format to binary preserves base64 encoding."""
+    run_main_and_assert(
+        input_path=OPEN_API_DATA_PATH / "type_mappings_byte_to_binary.yaml",
+        output_path=output_file,
+        input_file_type="openapi",
+        assert_func=assert_file_content,
+        expected_file="type_mappings_byte_to_binary.py",
+        extra_args=[
+            "--output-model-type",
+            "pydantic_v2.BaseModel",
+            "--type-mappings",
+            "string+byte=binary",
+        ],
+    )
+
+
 def test_main_openapi_unquoted_null(output_file: Path) -> None:
     """Test OpenAPI generation with unquoted null values."""
     run_main_and_assert(