Parse enum fixes. (#367)

Co-authored-by: Hasan Ramezani <[email protected]>

Author: kschwab

pydantic_settings/sources.py

@@ -678,9 +678,9 @@ def prepare_field_value(self, field_name: str, field: FieldInfo, value: Any, val
             ValuesError: When There is an error in deserializing value for complex field.
         """
         is_complex, allow_parse_failure = self._field_is_complex(field)
-        if self.env_parse_enums and lenient_issubclass(field.annotation, Enum):
-            if value in tuple(val.name for val in field.annotation):  # type: ignore
-                value = field.annotation[value]  # type: ignore
+        if self.env_parse_enums:
+            enum_val = _annotation_enum_name_to_val(field.annotation, value)
+            value = value if enum_val is None else enum_val
 
         if is_complex or value_is_complex:
             if isinstance(value, EnvNoneType):
@@ -1765,7 +1765,8 @@ def _help_format(self, field_name: str, field_info: FieldInfo, model_default: An
             elif model_default not in (PydanticUndefined, None) and callable(model_default):
                 default = f'(default factory: {self._metavar_format(model_default)})'
             elif field_info.default not in (PydanticUndefined, None):
-                default = f'(default: {field_info.default})'
+                enum_name = _annotation_enum_val_to_name(field_info.annotation, field_info.default)
+                default = f'(default: {field_info.default if enum_name is None else enum_name})'
             elif field_info.default_factory is not None:
                 default = f'(default: {field_info.default_factory})'
             _help += f' {default}' if _help else default
@@ -2075,3 +2076,19 @@ def _strip_annotated(annotation: Any) -> Any:
     while get_origin(annotation) == Annotated:
         annotation = get_args(annotation)[0]
     return annotation
+
+
+def _annotation_enum_val_to_name(annotation: type[Any] | None, value: Any) -> Optional[str]:
+    for type_ in (annotation, get_origin(annotation), *get_args(annotation)):
+        if lenient_issubclass(type_, Enum):
+            if value in tuple(val.value for val in type_):
+                return type_(value).name
+    return None
+
+
+def _annotation_enum_name_to_val(annotation: type[Any] | None, name: Any) -> Any:
+    for type_ in (annotation, get_origin(annotation), *get_args(annotation)):
+        if lenient_issubclass(type_, Enum):
+            if name in tuple(val.name for val in type_):
+                return type_[name]
+    return None

tests/test_settings.py

@@ -2116,9 +2116,11 @@ class Settings(BaseSettings):
 def test_env_parse_enums(env):
     class Settings(BaseSettings):
         fruit: FruitsEnum
+        union_fruit: Optional[Union[int, FruitsEnum]] = None
 
     with pytest.raises(ValidationError) as exc_info:
         env.set('FRUIT', 'kiwi')
+        env.set('UNION_FRUIT', 'kiwi')
         s = Settings()
     assert exc_info.value.errors(include_url=False) == [
         {
@@ -2127,18 +2129,42 @@ class Settings(BaseSettings):
             'msg': 'Input should be 0, 1 or 2',
             'input': 'kiwi',
             'ctx': {'expected': '0, 1 or 2'},
-        }
+        },
+        {
+            'input': 'kiwi',
+            'loc': (
+                'union_fruit',
+                'int',
+            ),
+            'msg': 'Input should be a valid integer, unable to parse string as an integer',
+            'type': 'int_parsing',
+        },
+        {
+            'ctx': {
+                'expected': '0, 1 or 2',
+            },
+            'input': 'kiwi',
+            'loc': (
+                'union_fruit',
+                'int-enum[FruitsEnum]',
+            ),
+            'msg': 'Input should be 0, 1 or 2',
+            'type': 'enum',
+        },
     ]
 
     env.set('FRUIT', str(FruitsEnum.lime.value))
+    env.set('UNION_FRUIT', str(FruitsEnum.lime.value))
     s = Settings()
     assert s.fruit == FruitsEnum.lime
 
     env.set('FRUIT', 'kiwi')
+    env.set('UNION_FRUIT', 'kiwi')
     s = Settings(_env_parse_enums=True)
     assert s.fruit == FruitsEnum.kiwi
 
     env.set('FRUIT', str(FruitsEnum.lime.value))
+    env.set('UNION_FRUIT', str(FruitsEnum.lime.value))
     s = Settings(_env_parse_enums=True)
     assert s.fruit == FruitsEnum.lime
 
@@ -3129,17 +3155,18 @@ class Cfg(BaseSettings):
     assert cfg.model_dump() == {'child': {'name': 'new name a', 'diff_a': 'new diff a'}}
 
 
-def test_cli_enums():
+def test_cli_enums(capsys, monkeypatch):
     class Pet(IntEnum):
         dog = 0
         cat = 1
         bird = 2
 
     class Cfg(BaseSettings):
-        pet: Pet
+        pet: Pet = Pet.dog
+        union_pet: Union[Pet, int] = 43
 
-    cfg = Cfg(_cli_parse_args=['--pet', 'cat'])
-    assert cfg.model_dump() == {'pet': Pet.cat}
+    cfg = Cfg(_cli_parse_args=['--pet', 'cat', '--union_pet', 'dog'])
+    assert cfg.model_dump() == {'pet': Pet.cat, 'union_pet': Pet.dog}
 
     with pytest.raises(ValidationError) as exc_info:
         Cfg(_cli_parse_args=['--pet', 'rock'])
@@ -3153,6 +3180,24 @@ class Cfg(BaseSettings):
         }
     ]
 
+    with monkeypatch.context() as m:
+        m.setattr(sys, 'argv', ['example.py', '--help'])
+
+        with pytest.raises(SystemExit):
+            Cfg(_cli_parse_args=True)
+        assert (
+            capsys.readouterr().out
+            == f"""usage: example.py [-h] [--pet {{dog,cat,bird}}]
+                  [--union_pet {{{{dog,cat,bird}},int}}]
+
+{ARGPARSE_OPTIONS_TEXT}:
+  -h, --help            show this help message and exit
+  --pet {{dog,cat,bird}}  (default: dog)
+  --union_pet {{{{dog,cat,bird}},int}}
+                        (default: 43)
+"""
+        )
+
 
 def test_cli_literals():
     class Cfg(BaseSettings):