Move enum validation and serialization to Rust (#9064)

Co-authored-by: sydney-runkle <sydneymarierunkle@gmail.com>
Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>

.gitignore

@@ -32,6 +32,7 @@ _build/
 /coverage.lcov
 .coverage
 test.py
+run_mypy_on_file.py
 
 # Documentation files
 /docs/changelog.md

docs/concepts/serialization.md

@@ -209,7 +209,7 @@ m = WithCustomEncoders(
     dt=datetime(2032, 6, 1, tzinfo=timezone.utc), diff=timedelta(hours=100)
 )
 print(m.model_dump_json())
-#> {"dt":1969660800.0,"diff":"P4DT14400S"}
+#> {"dt":1969660800.0,"diff":"P4DT4H"}
 
 
 class Model(BaseModel):

pydantic/_internal/_known_annotated_metadata.py

@@ -37,6 +37,7 @@
 TIMEDELTA_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *STRICT}
 TIME_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *STRICT}
 LAX_OR_STRICT_CONSTRAINTS = STRICT
+ENUM_CONSTRAINTS = STRICT
 
 UNION_CONSTRAINTS = {'union_mode'}
 URL_CONSTRAINTS = {
@@ -89,6 +90,8 @@
     CONSTRAINTS_TO_ALLOWED_SCHEMAS[constraint].update(('uuid',))
 for constraint in LAX_OR_STRICT_CONSTRAINTS:
     CONSTRAINTS_TO_ALLOWED_SCHEMAS[constraint].update(('lax-or-strict',))
+for constraint in ENUM_CONSTRAINTS:
+    CONSTRAINTS_TO_ALLOWED_SCHEMAS[constraint].update(('enum',))
 
 
 def add_js_update_schema(s: cs.CoreSchema, f: Callable[[], dict[str, Any]]) -> None:

pydantic/_internal/_std_types_schema.py

@@ -67,58 +67,65 @@ def get_enum_core_schema(enum_type: type[Enum], config: ConfigDict) -> CoreSchem
     js_updates = {'title': enum_type.__name__, 'description': description}
     js_updates = {k: v for k, v in js_updates.items() if v is not None}
 
-    def get_json_schema(_, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
-        json_schema = handler(core_schema.literal_schema([x.value for x in cases], ref=enum_ref))
-        original_schema = handler.resolve_ref_schema(json_schema)
-        update_json_schema(original_schema, js_updates)
-        return json_schema
-
-    if not cases:
-        # Use an isinstance check for enums with no cases.
-        # The most important use case for this is creating TypeVar bounds for generics that should
-        # be restricted to enums. This is more consistent than it might seem at first, since you can only
-        # subclass enum.Enum (or subclasses of enum.Enum) if all parent classes have no cases.
-        # We use the get_json_schema function when an Enum subclass has been declared with no cases
-        # so that we can still generate a valid json schema.
-        return core_schema.is_instance_schema(enum_type, metadata={'pydantic_js_functions': [get_json_schema]})
-
-    if len(cases) == 1:
-        expected = repr(cases[0].value)
-    else:
-        expected = ', '.join([repr(case.value) for case in cases[:-1]]) + f' or {cases[-1].value!r}'
-
-    def to_enum(input_value: Any, /) -> Enum:
-        try:
-            return enum_type(input_value)
-        except ValueError:
-            raise PydanticCustomError('enum', 'Input should be {expected}', {'expected': expected})
-
+    sub_type: typing_extensions.Literal['str', 'int', 'float'] | None = None
     if issubclass(enum_type, int):
-        # this handles `IntEnum`, and also `Foobar(int, Enum)`
-        js_updates['type'] = 'integer'
-        lax_schema = core_schema.no_info_after_validator_function(to_enum, core_schema.int_schema())
+        sub_type = 'int'
+        value_ser_type: core_schema.SerSchema = core_schema.simple_ser_schema('int')
     elif issubclass(enum_type, str):
         # this handles `StrEnum` (3.11 only), and also `Foobar(str, Enum)`
-        js_updates['type'] = 'string'
-        lax_schema = core_schema.no_info_after_validator_function(to_enum, core_schema.str_schema())
+        sub_type = 'str'
+        value_ser_type = core_schema.simple_ser_schema('str')
     elif issubclass(enum_type, float):
-        js_updates['type'] = 'numeric'
-        lax_schema = core_schema.no_info_after_validator_function(to_enum, core_schema.float_schema())
+        sub_type = 'float'
+        value_ser_type = core_schema.simple_ser_schema('float')
     else:
-        lax_schema = core_schema.no_info_plain_validator_function(to_enum)
+        # TODO this is an ugly hack, how do we trigger an Any schema for serialization?
+        value_ser_type = core_schema.plain_serializer_function_ser_schema(lambda x: x)
+
+    if cases:
+
+        def get_json_schema(schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
+            json_schema = handler(schema)
+            original_schema = handler.resolve_ref_schema(json_schema)
+            update_json_schema(original_schema, js_updates)
+            return json_schema
+
+        # we don't want to add the missing to the schema if it's the default one
+        default_missing = getattr(enum_type._missing_, '__func__', None) == Enum._missing_.__func__  # type: ignore
+        enum_schema = core_schema.enum_schema(
+            enum_type,
+            cases,
+            sub_type=sub_type,
+            missing=None if default_missing else enum_type._missing_,
+            ref=enum_ref,
+            metadata={'pydantic_js_functions': [get_json_schema]},
+        )
 
-    enum_schema = core_schema.lax_or_strict_schema(
-        lax_schema=lax_schema,
-        strict_schema=core_schema.json_or_python_schema(
-            json_schema=lax_schema, python_schema=core_schema.is_instance_schema(enum_type)
-        ),
-        ref=enum_ref,
-        metadata={'pydantic_js_functions': [get_json_schema]},
-    )
-    if config.get('use_enum_values', False):
-        enum_schema = core_schema.no_info_after_validator_function(attrgetter('value'), enum_schema)
+        if config.get('use_enum_values', False):
+            enum_schema = core_schema.no_info_after_validator_function(
+                attrgetter('value'), enum_schema, serialization=value_ser_type
+            )
 
-    return enum_schema
+        return enum_schema
+
+    else:
+
+        def get_json_schema_no_cases(_, handler: GetJsonSchemaHandler) -> JsonSchemaValue:
+            json_schema = handler(core_schema.enum_schema(enum_type, cases, sub_type=sub_type, ref=enum_ref))
+            original_schema = handler.resolve_ref_schema(json_schema)
+            update_json_schema(original_schema, js_updates)
+            return json_schema
+
+        # Use an isinstance check for enums with no cases.
+        # The most important use case for this is creating TypeVar bounds for generics that should
+        # be restricted to enums. This is more consistent than it might seem at first, since you can only
+        # subclass enum.Enum (or subclasses of enum.Enum) if all parent classes have no cases.
+        # We use the get_json_schema function when an Enum subclass has been declared with no cases
+        # so that we can still generate a valid json schema.
+        return core_schema.is_instance_schema(
+            enum_type,
+            metadata={'pydantic_js_functions': [get_json_schema_no_cases]},
+        )
 
 
 @dataclasses.dataclass(**slots_true)

pydantic/json_schema.py

@@ -742,13 +742,51 @@ def literal_schema(self, schema: core_schema.LiteralSchema) -> JsonSchemaValue:
         elif types == {int}:
             result['type'] = 'integer'
         elif types == {float}:
-            result['type'] = 'number'
+            result['type'] = 'numeric'
         elif types == {bool}:
             result['type'] = 'boolean'
         elif types == {list}:
             result['type'] = 'array'
         return result
 
+    def enum_schema(self, schema: core_schema.EnumSchema) -> JsonSchemaValue:
+        """Generates a JSON schema that matches an Enum value.
+
+        Args:
+            schema: The core schema.
+
+        Returns:
+            The generated JSON schema.
+        """
+        enum_type = schema['cls']
+        description = None if not enum_type.__doc__ else inspect.cleandoc(enum_type.__doc__)
+        if (
+            description == 'An enumeration.'
+        ):  # This is the default value provided by enum.EnumMeta.__new__; don't use it
+            description = None
+        result: dict[str, Any] = {'title': enum_type.__name__, 'description': description}
+        result = {k: v for k, v in result.items() if v is not None}
+
+        expected = [to_jsonable_python(v.value) for v in schema['members']]
+
+        result['enum'] = expected
+        if len(expected) == 1:
+            result['const'] = expected[0]
+
+        types = {type(e) for e in expected}
+        if isinstance(enum_type, str) or types == {str}:
+            result['type'] = 'string'
+        elif isinstance(enum_type, int) or types == {int}:
+            result['type'] = 'integer'
+        elif isinstance(enum_type, float) or types == {float}:
+            result['type'] = 'numeric'
+        elif types == {bool}:
+            result['type'] = 'boolean'
+        elif types == {list}:
+            result['type'] = 'array'
+
+        return result
+
     def is_instance_schema(self, schema: core_schema.IsInstanceSchema) -> JsonSchemaValue:
         """Handles JSON schema generation for a core schema that checks if a value is an instance of a class.
 

pyproject.toml

@@ -48,7 +48,7 @@ requires-python = '>=3.8'
 dependencies = [
     'typing-extensions>=4.6.1',
     'annotated-types>=0.4.0',
-    "pydantic-core==2.16.3",
+    "pydantic-core==2.17.0",
 ]
 dynamic = ['version', 'readme']
 

tests/test_fastapi.sh

@@ -18,4 +18,5 @@ cd .. && pip install . && cd fastapi
 # To skip a specific test, add '--deselect path/to/test.py::test_name' to the end of this command
 #
 # To update the list of deselected tests, remove all deselections, run the tests, and re-add any remaining failures
-./scripts/test.sh -vv
+# TODO remove this once that test is fixed, see https://github.com/pydantic/pydantic/pull/9064
+./scripts/test.sh -vv --deselect tests/test_tutorial/test_path_params/test_tutorial005.py::test_get_enums_invalid

tests/test_json.py

@@ -229,7 +229,7 @@ class Model(BaseModel):
 
     m = Model(x=123)
     json_data = m.model_dump_json()
-    assert json_data == '{"x":"PT123S"}'
+    assert json_data == '{"x":"PT2M3S"}'
     assert Model.model_validate_json(json_data).x == timedelta(seconds=123)
 
 
@@ -433,7 +433,7 @@ class Child(Parent):
         model_config = ConfigDict(json_encoders={datetime: lambda _: 'child_encoder'})
 
     # insert_assert(Child().model_dump())
-    assert json.loads(Child().model_dump_json()) == {'dt': 'child_encoder', 'timedt': 'P4DT14400S'}
+    assert json.loads(Child().model_dump_json()) == {'dt': 'child_encoder', 'timedt': 'P4DT4H'}
 
 
 def test_custom_iso_timedelta_annotated():

tests/test_json_schema.py

@@ -1728,7 +1728,7 @@ class Outer(BaseModel):
     'ser_json_timedelta,properties',
     [
         ('float', {'duration': {'default': 300.0, 'title': 'Duration', 'type': 'number'}}),
-        ('iso8601', {'duration': {'default': 'PT300S', 'format': 'duration', 'title': 'Duration', 'type': 'string'}}),
+        ('iso8601', {'duration': {'default': 'PT5M', 'format': 'duration', 'title': 'Duration', 'type': 'string'}}),
     ],
 )
 def test_model_default_timedelta(ser_json_timedelta: Literal['float', 'iso8601'], properties: typing.Dict[str, Any]):
@@ -1770,7 +1770,7 @@ class Model(BaseModel):
     'ser_json_timedelta,properties',
     [
         ('float', {'duration': {'default': 300.0, 'title': 'Duration', 'type': 'number'}}),
-        ('iso8601', {'duration': {'default': 'PT300S', 'format': 'duration', 'title': 'Duration', 'type': 'string'}}),
+        ('iso8601', {'duration': {'default': 'PT5M', 'format': 'duration', 'title': 'Duration', 'type': 'string'}}),
     ],
 )
 def test_dataclass_default_timedelta(
@@ -1812,7 +1812,7 @@ class Dataclass:
     'ser_json_timedelta,properties',
     [
         ('float', {'duration': {'default': 300.0, 'title': 'Duration', 'type': 'number'}}),
-        ('iso8601', {'duration': {'default': 'PT300S', 'format': 'duration', 'title': 'Duration', 'type': 'string'}}),
+        ('iso8601', {'duration': {'default': 'PT5M', 'format': 'duration', 'title': 'Duration', 'type': 'string'}}),
     ],
 )
 def test_typeddict_default_timedelta(

tests/test_types.py

@@ -1649,6 +1649,35 @@ class MyEnum(Enum):
     ]
 
 
+def test_enum_missing_default():
+    class MyEnum(Enum):
+        a = 1
+
+    ta = TypeAdapter(MyEnum)
+    missing_value = re.search(r'missing: (\w+)', repr(ta.validator)).group(1)
+    assert missing_value == 'None'
+
+    assert ta.validate_python(1) is MyEnum.a
+    with pytest.raises(ValidationError):
+        ta.validate_python(2)
+
+
+def test_enum_missing_custom():
+    class MyEnum(Enum):
+        a = 1
+
+        @classmethod
+        def _missing_(cls, value):
+            return MyEnum.a
+
+    ta = TypeAdapter(MyEnum)
+    missing_value = re.search(r'missing: (\w+)', repr(ta.validator)).group(1)
+    assert missing_value == 'Some'
+
+    assert ta.validate_python(1) is MyEnum.a
+    assert ta.validate_python(2) is MyEnum.a
+
+
 def test_int_enum_type():
     class Model(BaseModel):
         my_enum: IntEnum