refactor: serialization function

pydantic · Jan 12, 2024 · 90bfea5 · 90bfea5
1 parent 803766a
commit 90bfea5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 38 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
@@ -1793,28 +1795,6 @@ def _validate(cls, __input_value: Any, _: core_schema.ValidationInfo) -> ByteSiz
 
         return cls(int(float(scalar) * unit_mult))
 
-    @staticmethod
-    def _serialize(__input_value: Any) -> int:
-        try:
-            return int(__input_value)
-        except ValueError:
-            pass
-
-        str_match = byte_string_re.match(str(__input_value))
-        if str_match is None:
-            raise PydanticCustomError('byte_size', 'could not parse value and unit from byte string')
-
-        scalar, unit = str_match.groups()
-        if unit is None:
-            unit = 'b'
-
-        try:
-            unit_mult = BYTE_SIZES[unit.lower()]
-        except KeyError:
-            raise PydanticCustomError('byte_size_unit', 'could not interpret byte unit: {unit}', {'unit': unit})
-
-        return int(float(scalar) * unit_mult)
-
     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',