feat: Add basic typing support

.github/workflows/test.yml

@@ -10,7 +10,7 @@ concurrency:
 
 jobs:
   tests:
-    name: Python ${{ matrix.python-version }}, Database ${{ matrix.database-type }}
+    name: Python ${{ matrix.python-version }}, Env ${{ matrix.environment }}
     runs-on: ubuntu-latest
 
     strategy:
@@ -24,9 +24,10 @@ jobs:
           - "3.12"
           - "pypy-3.9"
           - "pypy-3.10"
-        database-type:
+        environment:
           - "sqlite"
           - "postgres"
+          - "typecheck"
 
     services:
       mongodb:
@@ -56,4 +57,4 @@ jobs:
       - name: Run tests
         run: tox
         env:
-          DATABASE_TYPE: ${{ matrix.database-type }}
+          ENVIRONMENT: ${{ matrix.environment }}

Makefile

@@ -66,6 +66,10 @@ test:
 		-Wdefault:"datetime.datetime.utcfromtimestamp() is deprecated and scheduled for removal in a future version.":DeprecationWarning:: \
 		-m unittest
 
+# Consider using pytest for entire test run (it just works)
+test-types:
+	pytest tests/typecheck --mypy-only-local-stub
+
 # DOC: Test the examples
 example-test:
 	$(MAKE) -C $(EXAMPLES_DIR) test

factory/base.py

@@ -4,11 +4,14 @@
 import collections
 import logging
 import warnings
+from typing import Generic, List, TypeVar
 
 from . import builder, declarations, enums, errors, utils
 
 logger = logging.getLogger('factory.generate')
 
+T = TypeVar('T')
+
 # Factory metaclasses
 
 
@@ -405,14 +408,14 @@ def reset(self, next_value=0):
         self.seq = next_value
 
 
-class BaseFactory:
+class BaseFactory(Generic[T]):
     """Factory base support for sequences, attributes and stubs."""
 
     # Backwards compatibility
     UnknownStrategy = errors.UnknownStrategy
     UnsupportedStrategy = errors.UnsupportedStrategy
 
-    def __new__(cls, *args, **kwargs):
+    def __new__(cls, *args, **kwargs) -> T:
         """Would be called if trying to instantiate the class."""
         raise errors.FactoryError('You cannot instantiate BaseFactory')
 
@@ -506,12 +509,12 @@ def _create(cls, model_class, *args, **kwargs):
         return model_class(*args, **kwargs)
 
     @classmethod
-    def build(cls, **kwargs):
+    def build(cls, **kwargs) -> T:
         """Build an instance of the associated class, with overridden attrs."""
         return cls._generate(enums.BUILD_STRATEGY, kwargs)
 
     @classmethod
-    def build_batch(cls, size, **kwargs):
+    def build_batch(cls, size: int, **kwargs) -> List[T]:
         """Build a batch of instances of the given class, with overridden attrs.
 
         Args:
@@ -523,12 +526,12 @@ def build_batch(cls, size, **kwargs):
         return [cls.build(**kwargs) for _ in range(size)]
 
     @classmethod
-    def create(cls, **kwargs):
+    def create(cls, **kwargs) -> T:
         """Create an instance of the associated class, with overridden attrs."""
         return cls._generate(enums.CREATE_STRATEGY, kwargs)
 
     @classmethod
-    def create_batch(cls, size, **kwargs):
+    def create_batch(cls, size: int, **kwargs) -> List[T]:
         """Create a batch of instances of the given class, with overridden attrs.
 
         Args:
@@ -627,7 +630,7 @@ def simple_generate_batch(cls, create, size, **kwargs):
         return cls.generate_batch(strategy, size, **kwargs)
 
 
-class Factory(BaseFactory, metaclass=FactoryMetaClass):
+class Factory(BaseFactory[T], metaclass=FactoryMetaClass):
     """Factory base with build and create support.
 
     This class has the ability to support multiple ORMs by using custom creation

factory/django.py

@@ -9,6 +9,7 @@
 import logging
 import os
 import warnings
+from typing import TypeVar
 
 from django.contrib.auth.hashers import make_password
 from django.core import files as django_files
@@ -20,7 +21,7 @@
 
 
 DEFAULT_DB_ALIAS = 'default'  # Same as django.db.DEFAULT_DB_ALIAS
-
+T = TypeVar("T")
 
 _LAZY_LOADS = {}
 
@@ -72,7 +73,7 @@ def get_model_class(self):
         return self.model
 
 
-class DjangoModelFactory(base.Factory):
+class DjangoModelFactory(base.Factory[T]):
     """Factory for Django models.
 
     This makes sure that the 'sequence' field of created objects is a new id.

setup.cfg

@@ -59,6 +59,10 @@ doc =
     sphinx_rtd_theme
     sphinxcontrib-spelling
 
+[options.package_data]
+factory =
+    py.typed
+
 [bdist_wheel]
 universal = 1
 

tests/typecheck/test_basics.yml

@@ -0,0 +1,13 @@
+- case: created_instance_is_of_correct_type
+  main: |
+    from django.db import models
+    import factory.django
+
+    class Book(models.Model):
+        ...
+
+    class BookFactory(factory.django.DjangoModelFactory[Book]):
+        ...
+
+    instance = BookFactory.create()
+    reveal_type(instance)  # N: Revealed type is "main.Book"

tox.ini

@@ -5,6 +5,7 @@ envlist =
     docs
     examples
     linkcheck
+    py{38,39,310,311,py39,py310}-typecheck
     py{38,39,310,311,312,py39,py310}-sqlite
     py{38,39,310,311,py39,py310}-django32-mongo-alchemy-{sqlite,postgres}
     py{38,39,310,311,py39,py310}-django41-mongo-alchemy-{sqlite,postgres}
@@ -25,7 +26,8 @@ python =
     pypy-3.10: pypy310
 
 [gh-actions:env]
-DATABASE_TYPE =
+ENVIRONMENT =
+    typecheck: typecheck
     sqlite: sqlite
     postgres: postgres
 
@@ -35,6 +37,9 @@ passenv =
     POSTGRES_HOST
     POSTGRES_DATABASE
 deps =
+    pytest
+    pytest-mypy-plugins
+    django-types
     alchemy: SQLAlchemy
     alchemy: sqlalchemy_utils
     mongo: mongoengine
@@ -81,3 +86,7 @@ extras = dev
 
 whitelist_externals = make
 commands = make lint
+
+[testenv:py{38,39,310,311,py39,py310}-typecheck]
+whitelist_externals = make
+commands = make test-types