diff --git a/polyfactory/value_generators/constrained_collections.py b/polyfactory/value_generators/constrained_collections.py index 3953313b..2638ef3d 100644 --- a/polyfactory/value_generators/constrained_collections.py +++ b/polyfactory/value_generators/constrained_collections.py @@ -1,5 +1,6 @@ from __future__ import annotations +from enum import EnumMeta from typing import TYPE_CHECKING, Any, Callable, List, Mapping, TypeVar, cast from polyfactory.exceptions import ParameterException @@ -39,6 +40,11 @@ def handle_constrained_collection( min_items = abs(min_items if min_items is not None else (max_items or 0)) max_items = abs(max_items if max_items is not None else min_items + 1) + if isinstance(field_meta.annotation, EnumMeta): + max_items = len(field_meta.annotation) + if min_items > max_items: + min_items = max_items + if max_items < min_items: msg = "max_items must be larger or equal to min_items" raise ParameterException(msg) diff --git a/tests/test_complex_types.py b/tests/test_complex_types.py index dcc699ff..ac5232cf 100644 --- a/tests/test_complex_types.py +++ b/tests/test_complex_types.py @@ -155,6 +155,25 @@ class MyFactory(ModelFactory): assert result.animal_list +def test_complex_typing_with_enum_set() -> None: + class Animal(str, Enum): + DOG = "Dog" + CAT = "Cat" + MONKEY = "Monkey" + + class MyModel(BaseModel): + animal_list: Set[Animal] + + class MyFactory(ModelFactory): + __model__ = MyModel + __randomize_collection_length__ = True + __min_collection_length__ = len(Animal) + 1 + __min_collection_length__ = len(Animal) + 2 + + result = MyFactory.build() + assert len(result.animal_list) == len(Animal) + + def test_union_literal() -> None: class MyModel(BaseModel): x: Union[int, Literal["a", "b", "c"]]