diff --git a/.gitignore b/.gitignore index f2ee4207..9c5f465c 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,7 @@ dist/ htmlcov/ MANIFEST tags + +# Virtualenv +.venv/ +venv/ diff --git a/factory/alchemy.py b/factory/alchemy.py index e782fbd8..00cdaa4a 100644 --- a/factory/alchemy.py +++ b/factory/alchemy.py @@ -43,7 +43,7 @@ def _build_default_options(self): ] -class SQLAlchemyModelFactory(base.Factory): +class SQLAlchemyModelFactory(base.Factory[base.T]): """Factory for SQLAlchemy models. """ _options_class = SQLAlchemyOptions diff --git a/factory/base.py b/factory/base.py index 454513be..65d1340c 100644 --- a/factory/base.py +++ b/factory/base.py @@ -666,7 +666,7 @@ def __init__(self, **kwargs): setattr(self, field, value) -class StubFactory(Factory): +class StubFactory(Factory[StubObject]): class Meta: strategy = enums.STUB_STRATEGY @@ -681,7 +681,7 @@ def create(cls, **kwargs): raise errors.UnsupportedStrategy() -class BaseDictFactory(Factory): +class BaseDictFactory(Factory[T]): """Factory for dictionary-like classes.""" class Meta: abstract = True @@ -698,12 +698,12 @@ def _create(cls, model_class, *args, **kwargs): return cls._build(model_class, *args, **kwargs) -class DictFactory(BaseDictFactory): +class DictFactory(BaseDictFactory[dict]): class Meta: model = dict -class BaseListFactory(Factory): +class BaseListFactory(Factory[T]): """Factory for list-like classes.""" class Meta: abstract = True @@ -724,7 +724,7 @@ def _create(cls, model_class, *args, **kwargs): return cls._build(model_class, *args, **kwargs) -class ListFactory(BaseListFactory): +class ListFactory(BaseListFactory[list]): class Meta: model = list diff --git a/factory/mogo.py b/factory/mogo.py index f886ae14..86bb50ac 100644 --- a/factory/mogo.py +++ b/factory/mogo.py @@ -7,7 +7,7 @@ from . import base -class MogoFactory(base.Factory): +class MogoFactory(base.Factory[base.T]): """Factory for mogo objects.""" class Meta: abstract = True diff --git a/factory/mongoengine.py b/factory/mongoengine.py index eb4a8dc5..0f462e9c 100644 --- a/factory/mongoengine.py +++ b/factory/mongoengine.py @@ -7,7 +7,7 @@ from . import base -class MongoEngineFactory(base.Factory): +class MongoEngineFactory(base.Factory[base.T]): """Factory for mongoengine objects.""" class Meta: diff --git a/tests/test_typing.py b/tests/test_typing.py index c2f8b564..dcf95573 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 @@ -14,9 +16,7 @@ class User: class TypingTests(unittest.TestCase): - def test_simple_factory(self) -> None: - class UserFactory(factory.Factory[User]): name = "John Doe" email = "john.doe@example.org" @@ -25,7 +25,20 @@ 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" + + self.assertTrue(isinstance(Pet.build(), dict)) + self.assertTrue(isinstance(Pet.create(), dict)) + + def test_list_factory(self) -> None: + self.assertTrue(isinstance(factory.ListFactory.build(), list)) + self.assertTrue(isinstance(factory.ListFactory.create(), list)) diff --git a/tox.ini b/tox.ini index 2bca8a08..158a6c3d 100644 --- a/tox.ini +++ b/tox.ini @@ -25,6 +25,7 @@ python = [testenv] deps = mypy + typing_extensions alchemy: SQLAlchemy mongo: mongoengine mongo: mongomock