Restore Django 5.0+ model_field_choices fixer (#542)

Fixes #417.

Restore the fixer with a new restriction that we only act when the
choices class is defined in the same file and can be reliably detected
as subclassing one of Django's enumeration types.

Co-authored-by: Thibaut Decombe <[email protected]>

Author: adamchainz

CHANGELOG.rst

@@ -6,6 +6,10 @@ Changelog
 
 * Add Django 5.2+ fixer ``staticfiles_find_all`` to rewrite calls to the staticfiles ``find()`` function using the old argument name ``all`` to the new name ``find_all``.
 
+* Restore Django 5.0+ fixer ``model_field_choices``, with a restriction to only apply when the enumeration type is defined in the same file.
+
+  Thanks to Thibaut Decombe for initially contributing it in `PR #369 <https://github.com/adamchainz/django-upgrade/pull/369>`__.
+
 * Add Django 2.0+ compatibility import rewrite of `django.http.cookie.SimpleCookie` to `http.cookies.SimpleCookie`.
 
   Thanks to Thibaut Decombe in `PR #537 <https://github.com/adamchainz/django-upgrade/pull/537>`__.

README.rst

@@ -333,6 +333,26 @@ Django 5.0
 
 `Release Notes <https://docs.djangoproject.com/en/5.0/releases/5.0/>`__
 
+Model field enumeration type ``.choices``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+**Name:** ``model_field_choices``
+
+Drop ``.choices`` for model field ``choices`` parameters, where the value is a class defined in the same file.
+This change is possible because the ``choices`` parameter now accepts enumeration types directly.
+
+.. code-block:: diff
+
+     from django.db import models
+
+     class Suit(models.IntegerChoices):
+         HEARTS = 1
+         ...
+
+     class Card(models.Model):
+    -    suit = models.IntegerField(choices=Suit.choices, default=Suit.DEFAULT)
+    +    suit = models.IntegerField(choices=Suit, default=Suit.DEFAULT)
+
 ``format_html()`` calls
 ~~~~~~~~~~~~~~~~~~~~~~~
 

src/django_upgrade/fixers/model_field_choices.py

@@ -0,0 +1,140 @@
+"""
+Drop `.choices` for model field `choices` parameters:
+https://docs.djangoproject.com/en/5.0/releases/5.0/#forms
+"""
+
+from __future__ import annotations
+
+import ast
+from collections import defaultdict
+from collections.abc import Iterable
+from functools import partial
+from typing import DefaultDict
+from typing import cast
+from weakref import WeakKeyDictionary
+
+from tokenize_rt import Offset
+from tokenize_rt import Token
+
+from django_upgrade.ast import ast_start_offset
+from django_upgrade.data import Fixer
+from django_upgrade.data import State
+from django_upgrade.data import TokenFunc
+from django_upgrade.tokens import OP
+from django_upgrade.tokens import find_last_token
+from django_upgrade.tokens import reverse_find
+
+fixer = Fixer(
+    __name__,
+    min_version=(5, 0),
+)
+
+# Cache defined enumeration types by module
+module_defined_enumeration_types: WeakKeyDictionary[ast.Module, dict[str, int]] = (
+    WeakKeyDictionary()
+)
+
+
+def defined_enumeration_types(module: ast.Module, up_to_line: int) -> set[str]:
+    """
+    Return a set of enumeration type class names defined in the given module, up to a line number.
+    """
+    if module not in module_defined_enumeration_types:
+        enum_dict = {}
+        from_imports: DefaultDict[str, set[str]] = defaultdict(set)
+        for node in module.body:
+            if (
+                isinstance(node, ast.ImportFrom)
+                and node.level == 0
+                and node.module is not None
+            ):
+                from_imports[node.module].update(
+                    name.name
+                    for name in node.names
+                    if name.asname is None and name.name != "*"
+                )
+            elif isinstance(node, ast.ClassDef):
+                # Check if the class inherits from one of Django's choice types
+                for base in node.bases:
+                    if _is_django_choices_type(from_imports, base):
+                        enum_dict[node.name] = node.lineno
+                        break
+        module_defined_enumeration_types[module] = enum_dict
+
+    return {
+        name
+        for name, line in module_defined_enumeration_types[module].items()
+        if line <= up_to_line
+    }
+
+
+DJANGO_CHOICES_TYPES = {
+    "TextChoices",
+    "IntegerChoices",
+    "Choices",
+}
+
+
+def _is_django_choices_type(
+    from_imports: DefaultDict[str, set[str]], node: ast.expr
+) -> bool:
+    """Check if an AST node refers to a Django enumeration type base class."""
+    return (
+        isinstance(node, ast.Name)
+        and node.id in DJANGO_CHOICES_TYPES
+        and (
+            node.id in from_imports["django.db.models"]
+            or node.id in from_imports["django.db.models.enums"]
+        )
+    ) or (
+        isinstance(node, ast.Attribute)
+        and node.attr in DJANGO_CHOICES_TYPES
+        and isinstance(node.value, ast.Name)
+        and (
+            (node.value.id == "models" and node.value.id in from_imports["django.db"])
+            or (
+                node.value.id == "enums"
+                and node.value.id in from_imports["django.db.models"]
+            )
+        )
+    )
+
+
+@fixer.register(ast.Call)
+def visit_Call(
+    state: State,
+    node: ast.Call,
+    parents: tuple[ast.AST, ...],
+) -> Iterable[tuple[Offset, TokenFunc]]:
+    if (
+        state.looks_like_models_file
+        and (
+            (
+                isinstance(node.func, ast.Attribute)
+                and isinstance(node.func.value, ast.Name)
+                and node.func.attr.endswith("Field")
+            )
+            or (isinstance(node.func, ast.Name) and node.func.id.endswith("Field"))
+        )
+        and any(
+            kw.arg == "choices"
+            and isinstance(kw.value, ast.Attribute)
+            and (target_node := kw.value).attr == "choices"
+            and isinstance(target_node.value, ast.Name)
+            and (
+                target_node.value.id
+                in defined_enumeration_types(
+                    cast(ast.Module, parents[0]),
+                    node.lineno,
+                )
+            )
+            for kw in node.keywords
+        )
+    ):
+        yield ast_start_offset(target_node), partial(remove_choices, node=target_node)
+
+
+def remove_choices(tokens: list[Token], i: int, node: ast.Attribute) -> None:
+    j = find_last_token(tokens, i, node=node)
+    i = reverse_find(tokens, j, name=OP, src=".")
+    del tokens[i : j + 1]