diff --git a/factory/base.py b/factory/base.py index f28f587c..46a43b9e 100644 --- a/factory/base.py +++ b/factory/base.py @@ -656,7 +656,7 @@ def __init__(self, **kwargs): setattr(self, field, value) -class StubFactory(Factory[T]): +class StubFactory(Factory[StubObject]): class Meta: strategy = enums.STUB_STRATEGY @@ -688,7 +688,7 @@ def _create(cls, model_class, *args, **kwargs): return cls._build(model_class, *args, **kwargs) -class DictFactory(BaseDictFactory[T]): +class DictFactory(BaseDictFactory[dict]): class Meta: model = dict @@ -714,7 +714,7 @@ def _create(cls, model_class, *args, **kwargs): return cls._build(model_class, *args, **kwargs) -class ListFactory(BaseListFactory[T]): +class ListFactory(BaseListFactory[list]): class Meta: model = list diff --git a/tests/test_typing.py b/tests/test_typing.py index c2f8b564..6ee98878 100644 --- a/tests/test_typing.py +++ b/tests/test_typing.py @@ -3,6 +3,8 @@ import dataclasses import unittest +from typing_extensions import assert_type + import factory @@ -25,7 +27,21 @@ class UserFactory(factory.Factory[User]): class Meta: model = User - result: User - result = UserFactory.build() - result = UserFactory.create() - self.assertEqual(result.name, "John Doe") + assert_type(UserFactory.build(), User) + assert_type(UserFactory.create(), User) + assert_type(UserFactory.build_batch(2), list[User]) + assert_type(UserFactory.create_batch(2), list[User]) + self.assertEqual(UserFactory.create().name, "John Doe") + + def test_dict_factory(self) -> None: + + class Pet(factory.DictFactory): + species = "dog" + name = "rover" + + assert_type(Pet.build(), dict) + assert_type(Pet.create(), dict) + + def test_list_factory(self) -> None: + assert_type(factory.ListFactory().build(), list) + assert_type(factory.ListFactory().create(), list) diff --git a/tox.ini b/tox.ini index d842c759..79e76f26 100644 --- a/tox.ini +++ b/tox.ini @@ -36,6 +36,7 @@ passenv = POSTGRES_DATABASE deps = mypy + typing_extensions alchemy: SQLAlchemy alchemy: sqlalchemy_utils mongo: mongoengine