fix: ensure to always return one of the values in Literal field type (

#2181)

* fix: ensure to always return one of the values in `Literal` field type

closes #2166

* perf: improve `literal_validator` speed

Thanks to @yobiscus

* fix: when more options in Literal

switch from `set` to `dict` to still have a O(1) complexity
Thanks @layday :)

changes/2166-PrettyWood.md

@@ -0,0 +1 @@
+fix: ensure to always return one of the values in `Literal` field type

pydantic/validators.py

@@ -442,12 +442,17 @@ def int_enum_validator(v: Any) -> IntEnum:
 
 def make_literal_validator(type_: Any) -> Callable[[Any], Any]:
     permitted_choices = all_literal_values(type_)
-    allowed_choices_set = set(permitted_choices)
+
+    # To have a O(1) complexity and still return one of the values set inside the `Literal`,
+    # we create a dict with the set values (a set causes some problems with the way intersection works).
+    # In some cases the set value and checked value can indeed be different (see `test_literal_validator_str_enum`)
+    allowed_choices = {v: v for v in permitted_choices}
 
     def literal_validator(v: Any) -> Any:
-        if v not in allowed_choices_set:
+        try:
+            return allowed_choices[v]
+        except KeyError:
             raise errors.WrongConstantError(given=v, permitted=permitted_choices)
-        return v
 
     return literal_validator
 

tests/test_validators.py

@@ -1,5 +1,6 @@
 from collections import deque
 from datetime import datetime
+from enum import Enum
 from itertools import product
 from typing import Dict, List, Optional, Tuple
 
@@ -1135,6 +1136,28 @@ class Model(BaseModel):
     ]
 
 
+@pytest.mark.skipif(not Literal, reason='typing_extensions not installed')
+def test_literal_validator_str_enum():
+    class Bar(str, Enum):
+        FIZ = 'fiz'
+        FUZ = 'fuz'
+
+    class Foo(BaseModel):
+        bar: Bar
+        barfiz: Literal[Bar.FIZ]
+        fizfuz: Literal[Bar.FIZ, Bar.FUZ]
+
+    my_foo = Foo.parse_obj({'bar': 'fiz', 'barfiz': 'fiz', 'fizfuz': 'fiz'})
+    assert my_foo.bar is Bar.FIZ
+    assert my_foo.barfiz is Bar.FIZ
+    assert my_foo.fizfuz is Bar.FIZ
+
+    my_foo = Foo.parse_obj({'bar': 'fiz', 'barfiz': 'fiz', 'fizfuz': 'fuz'})
+    assert my_foo.bar is Bar.FIZ
+    assert my_foo.barfiz is Bar.FIZ
+    assert my_foo.fizfuz is Bar.FUZ
+
+
 @pytest.mark.skipif(not Literal, reason='typing_extensions not installed')
 def test_nested_literal_validator():
     L1 = Literal['foo']