refactor: serialization function

pydantic · Jan 13, 2024 · 245a9f2 · 245a9f2
1 parent 15814e1
commit 245a9f2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 18 deletions.
diff --git a/pydantic/types.py b/pydantic/types.py
@@ -1768,7 +1768,9 @@ def __get_pydantic_core_schema__(cls, source: type[Any], handler: GetCoreSchemaH
                     core_schema.int_schema(ge=0),
                 ]
             ),
-            serialization=core_schema.plain_serializer_function_ser_schema(function=cls._serialize),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda v: int(v), return_schema=core_schema.int_schema(ge=0)
+            ),
         )
 
     @classmethod
@@ -1794,8 +1796,8 @@ def _validate(cls, __input_value: Any, _: core_schema.ValidationInfo) -> ByteSiz
         return cls(int(float(scalar) * unit_mult))
 
     @staticmethod
-    def _serialize(__input_value: Any) -> int:
-        return int(__input_value)
+    def _serialize(value: Any) -> int:
+        return int(value)
 
     def human_readable(self, decimal: bool = False) -> str:
         """Converts a byte size to a human readable string.

diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py
@@ -1316,6 +1316,8 @@ class Model(BaseModel):
     model_json_schema_validation = Model.model_json_schema(mode='validation')
     model_json_schema_serialization = Model.model_json_schema(mode='serialization')
 
+    print(model_json_schema_serialization)
+
     assert model_json_schema_validation == {
         'properties': {
             'a': {
@@ -1338,23 +1340,11 @@ class Model(BaseModel):
         'title': 'Model',
         'type': 'object',
     }
+
     assert model_json_schema_serialization == {
         'properties': {
-            'a': {
-                'anyOf': [
-                    {'pattern': '^\\s*(\\d*\\.?\\d+)\\s*(\\w+)?', 'type': 'string'},
-                    {'minimum': 0, 'type': 'integer'},
-                ],
-                'title': 'A',
-            },
-            'b': {
-                'anyOf': [
-                    {'pattern': '^\\s*(\\d*\\.?\\d+)\\s*(\\w+)?', 'type': 'string'},
-                    {'minimum': 0, 'type': 'integer'},
-                ],
-                'default': '1MB',
-                'title': 'B',
-            },
+            'a': {'minimum': 0, 'title': 'A', 'type': 'integer'},
+            'b': {'default': '1MB', 'minimum': 0, 'title': 'B', 'type': 'integer'},
         },
         'required': ['a'],
         'title': 'Model',