Add type annotations for primary key types by model [FC-0117]

src/openedx_catalog/api_impl.py

@@ -30,11 +30,11 @@ def get_catalog_course(*, org_code: str, course_code: str) -> CatalogCourse: ...
 @overload
 def get_catalog_course(*, key_str: str) -> CatalogCourse: ...
 @overload
-def get_catalog_course(*, pk: int) -> CatalogCourse: ...
+def get_catalog_course(*, pk: CatalogCourse.ID) -> CatalogCourse: ...
 
 
 def get_catalog_course(
-    pk: int | None = None,
+    pk: CatalogCourse.ID | None = None,
     key_str: str = "",
     org_code: str = "",
     course_code: str = "",
@@ -61,7 +61,7 @@ def get_catalog_course(
 
 
 def update_catalog_course(
-    catalog_course: CatalogCourse | int,
+    catalog_course: CatalogCourse | CatalogCourse.ID,
     *,
     title: str | None = None,  # Specify a string to change the title (display name).
     # The short language code (one of settings.ALL_LANGUAGES), e.g. "en", "es", "zh_HANS"
@@ -88,7 +88,7 @@ def update_catalog_course(
         cc.save(update_fields=update_fields)
 
 
-def delete_catalog_course(catalog_course: CatalogCourse | int) -> None:
+def delete_catalog_course(catalog_course: CatalogCourse | CatalogCourse.ID) -> None:
     """
     Delete a `CatalogCourse`. This will fail with a `ProtectedError` if any runs exist.
 

src/openedx_catalog/models/catalog_course.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+from typing import NewType
 
 from django.conf import settings
 from django.contrib import admin
@@ -12,7 +13,7 @@
 from django.utils.translation import gettext_lazy as _
 from organizations.models import Organization  # type: ignore[import]
 
-from openedx_django_lib.fields import case_insensitive_char_field, case_sensitive_char_field
+from openedx_django_lib.fields import TypedBigAutoField, case_insensitive_char_field, case_sensitive_char_field
 from openedx_django_lib.validators import validate_utc_datetime
 
 log = logging.getLogger(__name__)
@@ -59,7 +60,13 @@ class CatalogCourse(models.Model):
     courses in all instances of Open edX will need.)
     """
 
-    id = models.BigAutoField(
+    CatalogCourseID = NewType("CatalogCourseID", int)
+    type ID = CatalogCourseID
+
+    class IDField(TypedBigAutoField[ID]):  # Boilerplate for fully-typed ID field.
+        pass
+
+    id = IDField(
         primary_key=True,
         verbose_name=_("Primary Key"),
         help_text=_("The internal database ID for this catalog course. Should not be exposed to users nor in APIs."),

src/openedx_catalog/models/course_run.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+from typing import NewType
 
 from django.contrib import admin
 from django.core.exceptions import ValidationError
@@ -15,7 +16,7 @@
 from opaque_keys.edx.django.models import CourseKeyField
 from opaque_keys.edx.locator import CourseLocator
 
-from openedx_django_lib.fields import case_insensitive_char_field, case_sensitive_char_field
+from openedx_django_lib.fields import TypedBigAutoField, case_insensitive_char_field, case_sensitive_char_field
 from openedx_django_lib.validators import validate_utc_datetime
 
 from .catalog_course import CatalogCourse
@@ -79,8 +80,14 @@ class CourseRun(models.Model):
       learning package.
     """
 
+    CourseRunID = NewType("CourseRunID", int)
+    type ID = CourseRunID
+
+    class IDField(TypedBigAutoField[ID]):  # Boilerplate for fully-typed ID field.
+        pass
+
     # Use this field for relationships within the database:
-    id = models.BigAutoField(
+    id = IDField(
         primary_key=True,
         verbose_name=_("Primary Key"),
         help_text=_("The internal database ID for this course. Should not be exposed to users nor in APIs."),

src/openedx_content/applets/backup_restore/zipper.py

@@ -176,7 +176,7 @@ def get_publishable_entities(self) -> QuerySet[PublishableEntity]:
         Retrieve the publishable entities associated with the learning package.
         Prefetches related data for efficiency.
         """
-        lp_id = self.learning_package.pk
+        lp_id = self.learning_package.id
         publishable_entities: QuerySet[PublishableEntity] = publishing_api.get_publishable_entities(lp_id)
         return (
             publishable_entities  # type: ignore[no-redef]
@@ -210,7 +210,7 @@ def get_collections(self) -> QuerySet[Collection]:
         Get the collections associated with the learning package.
         """
         return (
-            collections_api.get_collections(self.learning_package.pk)
+            collections_api.get_collections(self.learning_package.id)
             .prefetch_related("entities")
         )
 
@@ -512,7 +512,7 @@ def __init__(self, zipf: zipfile.ZipFile, key: str | None = None, user: UserType
         self.user = user
         self.user_id = getattr(self.user, "id", None)
         self.lp_key = key  # If provided, use this key for the restored learning package
-        self.learning_package_id: int | None = None  # Will be set upon restoration
+        self.learning_package_id: LearningPackage.ID | None = None  # Will be set upon restoration
         self.utc_now: datetime = datetime.now(timezone.utc)
         self.component_types_cache: dict[tuple[str, str], ComponentType] = {}
         self.errors: list[dict[str, Any]] = []

src/openedx_content/applets/collections/api.py

@@ -10,7 +10,7 @@
 
 from ..publishing import api as publishing_api
 from ..publishing.models import PublishableEntity
-from .models import Collection, CollectionPublishableEntity
+from .models import Collection, CollectionPublishableEntity, LearningPackage
 
 # The public API that will be re-exported by openedx_content.api
 # is listed in the __all__ entries below. Internal helper functions that are
@@ -33,7 +33,7 @@
 
 
 def create_collection(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     key: str,
     *,
     title: str,
@@ -55,15 +55,15 @@ def create_collection(
     return collection
 
 
-def get_collection(learning_package_id: int, collection_key: str) -> Collection:
+def get_collection(learning_package_id: LearningPackage.ID, collection_key: str) -> Collection:
     """
     Get a Collection by ID
     """
     return Collection.objects.get_by_key(learning_package_id, collection_key)
 
 
 def update_collection(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     key: str,
     *,
     title: str | None = None,
@@ -89,7 +89,7 @@ def update_collection(
 
 
 def delete_collection(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     key: str,
     *,
     hard_delete=False,
@@ -111,7 +111,7 @@ def delete_collection(
 
 
 def restore_collection(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     key: str,
 ) -> Collection:
     """
@@ -125,7 +125,7 @@ def restore_collection(
 
 
 def add_to_collection(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     key: str,
     entities_qset: QuerySet[PublishableEntity],
     created_by: int | None = None,
@@ -145,7 +145,7 @@ def add_to_collection(
     invalid_entity = entities_qset.exclude(learning_package_id=learning_package_id).first()
     if invalid_entity:
         raise ValidationError(
-            f"Cannot add entity {invalid_entity.pk} in learning package {invalid_entity.learning_package_id} "
+            f"Cannot add entity {invalid_entity.id} in learning package {invalid_entity.learning_package_id} "
             f"to collection {key} in learning package {learning_package_id}."
         )
 
@@ -161,7 +161,7 @@ def add_to_collection(
 
 
 def remove_from_collection(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     key: str,
     entities_qset: QuerySet[PublishableEntity],
 ) -> Collection:
@@ -183,7 +183,7 @@ def remove_from_collection(
     return collection
 
 
-def get_entity_collections(learning_package_id: int, entity_key: str) -> QuerySet[Collection]:
+def get_entity_collections(learning_package_id: LearningPackage.ID, entity_key: str) -> QuerySet[Collection]:
     """
     Get all collections in the given learning package which contain this entity.
 
@@ -196,7 +196,10 @@ def get_entity_collections(learning_package_id: int, entity_key: str) -> QuerySe
     return entity.collections.filter(enabled=True).order_by("pk")
 
 
-def get_collection_entities(learning_package_id: int, collection_key: str) -> QuerySet[PublishableEntity]:
+def get_collection_entities(
+    learning_package_id: LearningPackage.ID,
+    collection_key: str,
+) -> QuerySet[PublishableEntity]:
     """
     Returns a QuerySet of PublishableEntities in a Collection.
 
@@ -208,7 +211,7 @@ def get_collection_entities(learning_package_id: int, collection_key: str) -> Qu
     ).order_by("pk")
 
 
-def get_collections(learning_package_id: int, enabled: bool | None = True) -> QuerySet[Collection]:
+def get_collections(learning_package_id: LearningPackage.ID, enabled: bool | None = True) -> QuerySet[Collection]:
     """
     Get all collections for a given learning package
 

src/openedx_content/applets/components/api.py

@@ -25,7 +25,7 @@
 
 from ..media import api as media_api
 from ..publishing import api as publishing_api
-from .models import Component, ComponentType, ComponentVersion, ComponentVersionMedia
+from .models import Component, ComponentType, ComponentVersion, ComponentVersionMedia, LearningPackage
 
 # The public API that will be re-exported by openedx_content.api
 # is listed in the __all__ entries below. Internal helper functions that are
@@ -96,7 +96,7 @@ def get_or_create_component_type_by_entity_key(entity_key: str) -> tuple[Compone
 
 
 def create_component(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     /,
     component_type: ComponentType,
     local_key: str,
@@ -127,7 +127,7 @@ def create_component(
 
 
 def create_component_version(
-    component_pk: int,
+    component_id: Component.ID,
     /,
     version_num: int,
     title: str,
@@ -139,21 +139,21 @@ def create_component_version(
     """
     with atomic():
         publishable_entity_version = publishing_api.create_publishable_entity_version(
-            component_pk,
+            component_id,
             version_num=version_num,
             title=title,
             created=created,
             created_by=created_by,
         )
         component_version = ComponentVersion.objects.create(
             publishable_entity_version=publishable_entity_version,
-            component_id=component_pk,
+            component_id=component_id,
         )
     return component_version
 
 
 def create_next_component_version(
-    component_pk: int,
+    component_id: Component.ID,
     /,
     media_to_replace: dict[str, int | None | bytes],
     created: datetime,
@@ -167,7 +167,7 @@ def create_next_component_version(
     Create a new ComponentVersion based on the most recent version.
 
     Args:
-        component_pk (int): The primary key of the Component to version.
+        component_id (int): The primary key of the Component to version.
         media_to_replace (dict): Mapping of file keys to Media IDs,
             None (for deletion), or bytes (for new file media).
         created (datetime): The creation timestamp for the new version.
@@ -218,7 +218,7 @@ def create_next_component_version(
     # should pick up from the last edited version. Likewise, a Draft might get
     # reverted to an earlier version, but we want the latest version_num when
     # creating the next version.
-    component = Component.objects.get(pk=component_pk)
+    component = Component.objects.get(pk=component_id)
     last_version = component.versioning.latest
     if last_version is None:
         next_version_num = 1
@@ -233,15 +233,15 @@ def create_next_component_version(
 
     with atomic():
         publishable_entity_version = publishing_api.create_publishable_entity_version(
-            component_pk,
+            component_id,
             version_num=next_version_num,
             title=title,
             created=created,
             created_by=created_by,
         )
         component_version = ComponentVersion.objects.create(
             publishable_entity_version=publishable_entity_version,
-            component_id=component_pk,
+            component_id=component_id,
         )
         # First copy the new stuff over...
         for key, media_pk_or_bytes in media_to_replace.items():
@@ -290,7 +290,7 @@ def create_next_component_version(
 
 
 def create_component_and_version(  # pylint: disable=too-many-positional-arguments
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     /,
     component_type: ComponentType,
     local_key: str,
@@ -313,7 +313,7 @@ def create_component_and_version(  # pylint: disable=too-many-positional-argumen
             can_stand_alone=can_stand_alone,
         )
         component_version = create_component_version(
-            component.pk,
+            component.id,
             version_num=1,
             title=title,
             created=created,
@@ -322,17 +322,17 @@ def create_component_and_version(  # pylint: disable=too-many-positional-argumen
         return (component, component_version)
 
 
-def get_component(component_pk: int, /) -> Component:
+def get_component(component_id: Component.ID, /) -> Component:
     """
     Get Component by its primary key.
 
     This is the same as the PublishableEntity's ID primary key.
     """
-    return Component.with_publishing_relations.get(pk=component_pk)
+    return Component.with_publishing_relations.get(pk=component_id)
 
 
 def get_component_by_key(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     /,
     namespace: str,
     type_name: str,
@@ -367,7 +367,7 @@ def get_component_version_by_uuid(uuid: UUID) -> ComponentVersion:
 
 
 def component_exists_by_key(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     /,
     namespace: str,
     type_name: str,
@@ -392,7 +392,7 @@ def component_exists_by_key(
 
 
 def get_components(  # pylint: disable=too-many-positional-arguments
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     /,
     draft: bool | None = None,
     published: bool | None = None,
@@ -435,7 +435,7 @@ def get_components(  # pylint: disable=too-many-positional-arguments
 
 
 def get_collection_components(
-    learning_package_id: int,
+    learning_package_id: LearningPackage.ID,
     collection_key: str,
 ) -> QuerySet[Component]:
     """