Add missing generic on Factory subclasses

factory/alchemy.py

@@ -1,10 +1,14 @@
 # Copyright: See the LICENSE file.
 
+from typing import TypeVar
+
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm.exc import NoResultFound
 
 from . import base, errors
 
+T = TypeVar("T")
+
 SESSION_PERSISTENCE_COMMIT = 'commit'
 SESSION_PERSISTENCE_FLUSH = 'flush'
 VALID_SESSION_PERSISTENCE_TYPES = [
@@ -46,7 +50,7 @@ def _build_default_options(self):
         ]
 
 
-class SQLAlchemyModelFactory(base.Factory):
+class SQLAlchemyModelFactory(base.Factory[T]):
     """Factory for SQLAlchemy models. """
 
     _options_class = SQLAlchemyOptions

factory/base.py

@@ -656,7 +656,7 @@ def __init__(self, **kwargs):
             setattr(self, field, value)
 
 
-class StubFactory(Factory):
+class StubFactory(Factory[T]):
 
     class Meta:
         strategy = enums.STUB_STRATEGY
@@ -671,7 +671,7 @@ def create(cls, **kwargs):
         raise errors.UnsupportedStrategy()
 
 
-class BaseDictFactory(Factory):
+class BaseDictFactory(Factory[T]):
     """Factory for dictionary-like classes."""
     class Meta:
         abstract = True
@@ -688,12 +688,12 @@ def _create(cls, model_class, *args, **kwargs):
         return cls._build(model_class, *args, **kwargs)
 
 
-class DictFactory(BaseDictFactory):
+class DictFactory(BaseDictFactory[T]):
     class Meta:
         model = dict
 
 
-class BaseListFactory(Factory):
+class BaseListFactory(Factory[T]):
     """Factory for list-like classes."""
     class Meta:
         abstract = True
@@ -714,7 +714,7 @@ def _create(cls, model_class, *args, **kwargs):
         return cls._build(model_class, *args, **kwargs)
 
 
-class ListFactory(BaseListFactory):
+class ListFactory(BaseListFactory[T]):
     class Meta:
         model = list
 

factory/mogo.py

@@ -2,12 +2,13 @@
 
 
 """factory_boy extensions for use with the mogo library (pymongo wrapper)."""
-
+from typing import TypeVar
 
 from . import base
 
+T = TypeVar("T")
 
-class MogoFactory(base.Factory):
+class MogoFactory(base.Factory[T]):
     """Factory for mogo objects."""
     class Meta:
         abstract = True

factory/mongoengine.py

@@ -2,12 +2,13 @@
 
 
 """factory_boy extensions for use with the mongoengine library (pymongo wrapper)."""
-
+from typing import TypeVar
 
 from . import base
 
+T = TypeVar("T")
 
-class MongoEngineFactory(base.Factory):
+class MongoEngineFactory(base.Factory[T]):
     """Factory for mongoengine objects."""
 
     class Meta: