refactor: improve typing of mixins

Author: bradenmacdonald

openedx_learning/apps/authoring/components/models.py

@@ -17,6 +17,8 @@
 """
 from __future__ import annotations
 
+from typing import ClassVar
+
 from django.db import models
 
 from ....lib.fields import case_sensitive_char_field, immutable_uuid_field, key_field
@@ -76,7 +78,7 @@ def __str__(self):
         return f"{self.namespace}:{self.name}"
 
 
-class Component(PublishableEntityMixin):  # type: ignore[django-manager-missing]
+class Component(PublishableEntityMixin):
     """
     This represents any Component that has ever existed in a LearningPackage.
 
@@ -120,14 +122,12 @@ class Component(PublishableEntityMixin):  # type: ignore[django-manager-missing]
     Make a foreign key to the Component model when you need a stable reference
     that will exist for as long as the LearningPackage itself exists.
     """
-    # Tell mypy what type our objects manager has.
-    # It's actually PublishableEntityMixinManager, but that has the exact same
-    # interface as the base manager class.
-    objects: models.Manager[Component] = WithRelationsManager(
+    # Set up our custom manager. It has the same API as the default one, but selects related objects by default.
+    objects: ClassVar[WithRelationsManager[Component]] = WithRelationsManager(  # type: ignore[assignment]
         'component_type'
     )
 
-    with_publishing_relations: models.Manager[Component] = WithRelationsManager(
+    with_publishing_relations = WithRelationsManager(
         'component_type',
         'publishable_entity',
         'publishable_entity__draft__version',
@@ -201,10 +201,6 @@ class ComponentVersion(PublishableEntityVersionMixin):
     This holds the content using a M:M relationship with Content via
     ComponentVersionContent.
     """
-    # Tell mypy what type our objects manager has.
-    # It's actually PublishableEntityVersionMixinManager, but that has the exact
-    # same interface as the base manager class.
-    objects: models.Manager[ComponentVersion]
 
     # This is technically redundant, since we can get this through
     # publishable_entity_version.publishable.component, but this is more

openedx_learning/apps/authoring/publishing/model_mixins.py

@@ -4,10 +4,12 @@
 from __future__ import annotations
 
 from functools import cached_property
+from typing import ClassVar, Self
 
 from django.core.exceptions import ImproperlyConfigured
 from django.db import models
-from django.db.models.query import QuerySet
+
+from openedx_learning.lib.managers import WithRelationsManager
 
 from .models import PublishableEntity, PublishableEntityVersion
 
@@ -28,17 +30,12 @@ class PublishableEntityMixin(models.Model):
     the publishing app's api.register_content_models (see its docstring for
     details).
     """
-
-    class PublishableEntityMixinManager(models.Manager):
-        def get_queryset(self) -> QuerySet:
-            return super().get_queryset() \
-                          .select_related(
-                              "publishable_entity",
-                              "publishable_entity__published",
-                              "publishable_entity__draft",
-                          )
-
-    objects: models.Manager[PublishableEntityMixin] = PublishableEntityMixinManager()
+    # select these related entities by default for all queries
+    objects: ClassVar[WithRelationsManager[Self]] = WithRelationsManager(
+        "publishable_entity",
+        "publishable_entity__published",
+        "publishable_entity__draft",
+    )
 
     publishable_entity = models.OneToOneField(
         PublishableEntity, on_delete=models.CASCADE, primary_key=True
@@ -294,17 +291,10 @@ class PublishableEntityVersionMixin(models.Model):
     details).
     """
 
-    class PublishableEntityVersionMixinManager(models.Manager):
-        def get_queryset(self) -> QuerySet:
-            return (
-                super()
-                .get_queryset()
-                .select_related(
-                    "publishable_entity_version",
-                )
-            )
-
-    objects: models.Manager[PublishableEntityVersionMixin] = PublishableEntityVersionMixinManager()
+    # select these related entities by default for all queries
+    objects: ClassVar[WithRelationsManager[Self]] = WithRelationsManager(
+        "publishable_entity_version",
+    )
 
     publishable_entity_version = models.OneToOneField(
         PublishableEntityVersion, on_delete=models.CASCADE, primary_key=True

openedx_learning/lib/managers.py

@@ -1,11 +1,17 @@
 """
 Custom Django ORM Managers.
 """
+from __future__ import annotations
+
+from typing import TypeVar
+
 from django.db import models
 from django.db.models.query import QuerySet
 
+M = TypeVar('M', bound=models.Model)
+
 
-class WithRelationsManager(models.Manager):
+class WithRelationsManager(models.Manager[M]):
     """
     Custom Manager that adds select_related to the default queryset.
 

tests/openedx_learning/apps/authoring/components/test_models.py

@@ -2,6 +2,7 @@
 Tests related to the Component models
 """
 from datetime import datetime, timezone
+from typing import TYPE_CHECKING, assert_type
 
 from freezegun import freeze_time
 
@@ -10,7 +11,7 @@
     get_component,
     get_or_create_component_type,
 )
-from openedx_learning.apps.authoring.components.models import ComponentType
+from openedx_learning.apps.authoring.components.models import Component, ComponentType, ComponentVersion
 from openedx_learning.apps.authoring.publishing.api import (
     LearningPackage,
     create_learning_package,
@@ -19,6 +20,15 @@
 )
 from openedx_learning.lib.test_utils import TestCase
 
+if TYPE_CHECKING:
+    # Test that our mixins on Component.objects and PublishableEntityVersionMixin etc. haven't broken manager typing
+    assert_type(Component.objects.create(), Component)
+    assert_type(Component.objects.get(), Component)
+    assert_type(Component.with_publishing_relations.create(), Component)
+    assert_type(Component.with_publishing_relations.get(), Component)
+    assert_type(ComponentVersion.objects.create(), ComponentVersion)
+    assert_type(ComponentVersion.objects.get(), ComponentVersion)
+
 
 class TestModelVersioningQueries(TestCase):
     """

tests/openedx_learning/apps/authoring/publishing/test_model_mixins.py

@@ -0,0 +1,24 @@
+"""
+Tests related to the Component models
+"""
+from typing import TYPE_CHECKING, assert_type
+
+from openedx_learning.apps.authoring.publishing.model_mixins import (
+    PublishableEntityMixin,
+    PublishableEntityVersionMixin,
+)
+from openedx_learning.lib.managers import WithRelationsManager
+
+if TYPE_CHECKING:
+    # Test that our mixins provide the right typing for 'objects'
+    class FooEntity(PublishableEntityMixin):
+        pass
+
+    assert_type(FooEntity.objects.create(), FooEntity)
+    assert_type(FooEntity.objects, WithRelationsManager[FooEntity])
+
+    class FooEntityVersion(PublishableEntityVersionMixin):
+        pass
+
+    assert_type(FooEntityVersion.objects.create(), FooEntityVersion)
+    assert_type(FooEntityVersion.objects, WithRelationsManager[FooEntityVersion])