Validate select_related lookups (#2806)

Author: UnknownPlatypus

mypy_django_plugin/main.py

@@ -168,6 +168,7 @@ def manager_and_queryset_method_hooks(self) -> dict[str, Callable[[MethodContext
             "prefetch_related": partial(
                 querysets.extract_prefetch_related_annotations, django_context=self.django_context
             ),
+            "select_related": partial(querysets.validate_select_related, django_context=self.django_context),
         }
 
     def get_method_hook(self, fullname: str) -> Callable[[MethodContext], MypyType] | None:

mypy_django_plugin/transformers/querysets.py

@@ -619,3 +619,79 @@ def extract_prefetch_related_annotations(ctx: MethodContext, django_context: Dja
         row_type = annotated_model
 
     return default_return_type.copy_modified(args=[annotated_model, row_type])
+
+
+def _get_select_related_field_choices(model_cls: type[Model]) -> set[str]:
+    """
+    Get valid field choices for select_related lookups.
+    Based on Django's SQLCompiler.get_related_selections._get_field_choices method.
+    """
+    opts = model_cls._meta
+
+    # Direct relation fields (forward relations)
+    direct_choices = (f.name for f in opts.fields if f.is_relation)
+
+    # Reverse relation fields (backward relations with unique=True)
+    reverse_choices = (f.field.related_query_name() for f in opts.related_objects if f.field.unique)
+    return {*direct_choices, *reverse_choices}
+
+
+def _validate_select_related_lookup(
+    ctx: MethodContext,
+    django_context: DjangoContext,
+    model_cls: type[Model],
+    lookup: str,
+) -> bool:
+    """Validate a single select_related lookup string."""
+    if not lookup.strip():
+        ctx.api.fail(
+            f'Invalid field name "{lookup}" in select_related lookup',
+            ctx.context,
+        )
+        return False
+
+    lookup_parts = lookup.split("__")
+    observed_model = model_cls
+    for i, part in enumerate(lookup_parts):
+        valid_choices = _get_select_related_field_choices(observed_model)
+
+        if part not in valid_choices:
+            ctx.api.fail(
+                f'Invalid field name "{part}" in select_related lookup. '
+                f"Choices are: {', '.join(sorted(valid_choices)) or '(none)'}",
+                ctx.context,
+            )
+            return False
+
+        if i < len(lookup_parts) - 1:  # Not the last part
+            try:
+                field, observed_model = django_context.resolve_lookup_into_field(observed_model, part)
+                if field is None:
+                    return False
+            except (FieldError, LookupsAreUnsupported):
+                # For good measure, but we should never reach this since we already validated the part name
+                return False
+
+    return True
+
+
+def validate_select_related(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
+    """
+    Validates that all lookup strings passed to select_related() resolve to actual model fields and relations.
+
+    Extracted and adapted from `django.db.models.sql.compiler.SQLCompiler.get_related_selections`
+    """
+    if not (
+        isinstance(ctx.type, Instance)
+        and (django_model := helpers.get_model_info_from_qs_ctx(ctx, django_context)) is not None
+        and ctx.arg_types
+        and ctx.arg_types[0]
+    ):
+        return ctx.default_return_type
+
+    for lookup_type in ctx.arg_types[0]:
+        lookup_value = helpers.get_literal_str_type(get_proper_type(lookup_type))
+        if lookup_value is not None:
+            _validate_select_related_lookup(ctx, django_context, django_model.cls, lookup_value)
+
+    return ctx.default_return_type

tests/typecheck/managers/querysets/test_basic_methods.yml

@@ -99,8 +99,8 @@
         Book.objects.all().select_related()
         Book.objects.select_related(None)
         Book.objects.all().select_related(None)
-        Book.objects.select_related("author", "dd") # Should be Error on invalid lookup 'dd' !
-        Book.objects.all().select_related("author", "dd") # Should be Error on invalid lookup 'dd' !
+        Book.objects.select_related("author", "dd") # E: Invalid field name "dd" in select_related lookup. Choices are: author  [misc]
+        Book.objects.all().select_related("author", "dd") # E: Invalid field name "dd" in select_related lookup. Choices are: author  [misc]
 
         Book.objects.select_related("None", None)
         Book.objects.all().select_related("None", None)

tests/typecheck/managers/querysets/test_select_related.yml

@@ -0,0 +1,137 @@
+-   case: select_related_valid_lookups
+    installed_apps:
+        - myapp
+    main: |
+        from myapp.models import Article, Author, Category, Blog, Post
+
+        # Valid forward relations
+        Article.objects.select_related("author")
+        Article.objects.select_related("category")
+
+        # Valid reverse relations (unique=True)
+        Author.objects.select_related("profile")
+
+        # Valid chained lookups
+        Article.objects.select_related("author__profile")
+        Article.objects.select_related("category__parent")
+        Post.objects.select_related("blog__owner")
+        Post.objects.select_related("category__parent__parent")
+
+        # Multiple valid lookups
+        Article.objects.select_related("author", "category")
+        Article.objects.select_related("author__profile", "category")
+
+        # Self-referential relationships
+        Category.objects.select_related("parent")
+        Category.objects.select_related("parent__parent")
+
+        # Variables containing strings
+        author_lookup = "aaa"
+        Article.objects.select_related(author_lookup)
+
+        # Dynamic lookups (should not be validated)
+        def get_lookup() -> str:
+            return "author"
+
+        Article.objects.select_related(get_lookup())
+
+        # Chaining select_related calls
+        qs1 = Article.objects.select_related("author")
+        qs2 = qs1.select_related("category")
+        Article.objects.select_related("author").select_related("category")
+
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class Category(models.Model):
+                    parent = models.ForeignKey('self', on_delete=models.CASCADE, null=True, blank=True)
+
+                class Author(models.Model):
+                    pass
+
+                class AuthorProfile(models.Model):
+                    author = models.OneToOneField(Author, on_delete=models.CASCADE, related_name='profile')
+
+                class Blog(models.Model):
+                    owner = models.ForeignKey(Author, on_delete=models.CASCADE)
+
+                class Article(models.Model):
+                    title = models.CharField(max_length=200)  # Keep for testing non-relation field
+                    author = models.ForeignKey(Author, on_delete=models.CASCADE)
+                    category = models.ForeignKey(Category, on_delete=models.CASCADE)
+
+                class Post(models.Model):
+                    blog = models.ForeignKey(Blog, on_delete=models.CASCADE)
+                    author = models.ForeignKey(Author, on_delete=models.CASCADE)
+                    category = models.ForeignKey(Category, on_delete=models.CASCADE)
+
+-   case: select_related_invalid_lookups
+    installed_apps:
+        - myapp
+    main: |
+        from myapp.models import Article, Author, Category, Post
+        from typing import Literal
+
+        # Invalid field names
+        Article.objects.select_related("nonexistent")  # E: Invalid field name "nonexistent" in select_related lookup. Choices are: author, category  [misc]
+        Article.objects.select_related("title")  # E: Invalid field name "title" in select_related lookup. Choices are: author, category  [misc]
+
+        # Invalid chained lookups
+        Article.objects.select_related("author__nonexistent")  # E: Invalid field name "nonexistent" in select_related lookup. Choices are: profile  [misc]
+        Article.objects.select_related("category__invalid")  # E: Invalid field name "invalid" in select_related lookup. Choices are: parent  [misc]
+
+        # Reverse many-to-many (not allowed)
+        Category.objects.select_related("article_set")  # E: Invalid field name "article_set" in select_related lookup. Choices are: parent  [misc]
+
+        # Mixed valid and invalid in same call
+        Post.objects.select_related("blog", "invalid")  # E: Invalid field name "invalid" in select_related lookup. Choices are: author, blog, category  [misc]
+
+        # Intermediary variable not validated if non literal
+        invalid_lookup = "nonexistent"
+        Article.objects.select_related(invalid_lookup)  # No error - variables not validated
+
+        # Intermediary variable validated if literal
+        invalid_lookup2: Literal["nonexistent"] = "nonexistent"
+        Article.objects.select_related(invalid_lookup2)  # E: Invalid field name "nonexistent" in select_related lookup. Choices are: author, category  [misc]
+
+        # Chaining with invalid lookups
+        qs3 = Article.objects.select_related("author").select_related("invalid")  # E: Invalid field name "invalid" in select_related lookup. Choices are: author, category  [misc]
+        Article.objects.select_related("author").select_related("invalid")  # E: Invalid field name "invalid" in select_related lookup. Choices are: author, category  [misc]
+
+        # Multiple arguments with invalid lookups
+        Article.objects.select_related("author", "invalid", "category")  # E: Invalid field name "invalid" in select_related lookup. Choices are: author, category  [misc]
+
+        # Whitespace-only lookup (invalid)
+        Article.objects.select_related("")   # E: Invalid field name "" in select_related lookup  [misc]
+        Article.objects.select_related("   ")  # E: Invalid field name "   " in select_related lookup  [misc]
+
+    files:
+        -   path: myapp/__init__.py
+        -   path: myapp/models.py
+            content: |
+                from django.db import models
+
+                class Category(models.Model):
+                    parent = models.ForeignKey('self', on_delete=models.CASCADE, null=True, blank=True)
+
+                class Author(models.Model):
+                    pass
+
+                class AuthorProfile(models.Model):
+                    author = models.OneToOneField(Author, on_delete=models.CASCADE, related_name='profile')
+
+                class Blog(models.Model):
+                    owner = models.ForeignKey(Author, on_delete=models.CASCADE)
+
+                class Article(models.Model):
+                    title = models.CharField(max_length=200)  # Keep for testing non-relation field
+                    author = models.ForeignKey(Author, on_delete=models.CASCADE)
+                    category = models.ForeignKey(Category, on_delete=models.CASCADE)
+
+                class Post(models.Model):
+                    blog = models.ForeignKey(Blog, on_delete=models.CASCADE)
+                    author = models.ForeignKey(Author, on_delete=models.CASCADE)
+                    category = models.ForeignKey(Category, on_delete=models.CASCADE)