fix: flatten_annotation behaviour for Optional (#440)

Author: sam-or

polyfactory/utils/helpers.py

@@ -94,8 +94,8 @@ def flatten_annotation(annotation: Any) -> list[Any]:
     if is_new_type(annotation):
         flat.extend(flatten_annotation(unwrap_new_type(annotation)))
     elif is_optional(annotation):
-        flat.append(NoneType)
-        flat.extend(flatten_annotation(arg) for arg in get_args(annotation) if arg not in (NoneType, None))
+        for a in get_args(annotation):
+            flat.extend(flatten_annotation(a))
     elif is_annotated(annotation):
         flat.extend(flatten_annotation(get_args(annotation)[0]))
     elif is_union(annotation):

tests/test_type_coverage_generation.py

@@ -3,7 +3,7 @@
 
 from dataclasses import dataclass, make_dataclass
 from datetime import date
-from typing import Dict, FrozenSet, List, Literal, Set, Tuple, Union
+from typing import Dict, FrozenSet, List, Literal, Optional, Set, Tuple, Union
 from uuid import UUID
 
 import pytest
@@ -212,3 +212,18 @@ class Factory(DataclassFactory[Model]):
 
     with pytest.raises(ParameterException):
         list(Factory.coverage())
+
+
+def test_coverage_optional_field() -> None:
+    @dataclass
+    class OptionalInt:
+        i: Optional[int]
+
+    class OptionalIntFactory(DataclassFactory[OptionalInt]):
+        __model__ = OptionalInt
+
+    results = list(OptionalIntFactory.coverage())
+    assert len(results) == 2
+
+    assert isinstance(results[0].i, int)
+    assert results[1].i is None