Add enum provider (#1729)

faker/documentor.py

@@ -1,13 +1,21 @@
 import inspect
 import warnings
 
-from typing import Any, Dict, List, Optional, Tuple, Union
+from enum import Enum, auto
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 from .generator import Generator
 from .providers import BaseProvider
 from .proxy import Faker
 
 
+class FakerEnum(Enum):
+    """Required for faker.providers.enum"""
+
+    A = auto
+    B = auto
+
+
 class Documentor:
     def __init__(self, generator: Union[Generator, Faker]) -> None:
         """
@@ -52,7 +60,7 @@ def get_provider_formatters(
                 continue
 
             arguments = []
-            faker_args: List[str] = []
+            faker_args: List[Union[str, Type[Enum]]] = []
             faker_kwargs = {}
 
             if name == "binary":
@@ -65,6 +73,9 @@ def get_provider_formatters(
                     }
                 )
 
+            if name == "enum":
+                faker_args = [FakerEnum]
+
             if with_args:
                 # retrieve all parameter
                 argspec = inspect.getfullargspec(method)

faker/providers/python/__init__.py

@@ -4,12 +4,19 @@
 import warnings
 
 from decimal import Decimal
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Type, Union, no_type_check
+from enum import Enum
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Type, TypeVar, Union, cast, no_type_check
 
+from ...exceptions import BaseFakerException
 from .. import BaseProvider, ElementsType
 
 TypesNames = List[str]
 TypesSpec = Union[List[Type], Tuple[Type, ...]]
+TEnum = TypeVar("TEnum", bound=Enum)
+
+
+class EmptyEnumException(BaseFakerException):
+    pass
 
 
 class Provider(BaseProvider):
@@ -417,3 +424,24 @@ def pystruct(
                 },
             }
         return types, d, nd
+
+    def enum(self, enum_cls: Type[TEnum]) -> TEnum:
+        """
+        Returns a random enum of the provided input `Enum` type.
+
+        :param enum_cls: The `Enum` type to produce the value for.
+        :returns: A randomly selected enum value.
+        """
+
+        if enum_cls is None:
+            raise ValueError("'enum_cls' cannot be None")
+
+        if not issubclass(enum_cls, Enum):
+            raise TypeError("'enum_cls' must be an Enum type")
+
+        members: List[TEnum] = list(cast(Iterable[TEnum], enum_cls))
+
+        if len(members) < 1:
+            raise EmptyEnumException(f"The provided Enum: '{enum_cls.__name__}' has no members.")
+
+        return self.random_element(members)

tests/providers/test_enum.py

@@ -0,0 +1,51 @@
+from enum import Enum, auto
+
+import pytest
+
+from faker.providers.python import EmptyEnumException
+
+
+class _TestEnumWithNoElements(Enum):
+    pass
+
+
+class _TestEnumWithSingleElement(Enum):
+    Single = auto
+
+
+class _TestEnum(Enum):
+    A = auto
+    B = auto
+    C = auto
+
+
+class TestEnumProvider:
+
+    num_samples = 100
+
+    def test_enum(self, faker, num_samples):
+        # (1/3) ** 100 ~ 1.94e-48 probability of this test failing because a specific
+        # value was not sampled
+        for _ in range(num_samples):
+            actual = faker.enum(_TestEnum)
+            assert actual in (_TestEnum.A, _TestEnum.B, _TestEnum.C)
+
+    def test_enum_single(self, faker):
+        assert faker.enum(_TestEnumWithSingleElement) == _TestEnumWithSingleElement.Single
+        assert faker.enum(_TestEnumWithSingleElement) == _TestEnumWithSingleElement.Single
+
+    def test_empty_enum_raises(self, faker):
+        with pytest.raises(
+            EmptyEnumException,
+            match="The provided Enum: '_TestEnumWithNoElements' has no members.",
+        ):
+            faker.enum(_TestEnumWithNoElements)
+
+    def test_none_raises(self, faker):
+        with pytest.raises(ValueError):
+            faker.enum(None)
+
+    def test_incorrect_type_raises(self, faker):
+        not_an_enum_type = type("NotAnEnumType")
+        with pytest.raises(TypeError):
+            faker.enum(not_an_enum_type)