feat: Add generic typing to factories

.gitignore

@@ -16,3 +16,7 @@ dist/
 htmlcov/
 MANIFEST
 tags
+
+# Virtualenv
+.venv/
+venv/

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

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
 

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

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:

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 = "[email protected]"
@@ -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))

tox.ini

@@ -25,6 +25,7 @@ python =
 [testenv]
 deps =
     mypy
+    typing_extensions
     alchemy: SQLAlchemy
     mongo: mongoengine
     mongo: mongomock