fix referencing .value of a Choices variable (#32)

Co-authored-by: Anthony Sottile <[email protected]>

Author: asottile-sentry

mypy_django_plugin/transformers/choices.py

@@ -1,4 +1,4 @@
-from mypy.nodes import MemberExpr, NameExpr, TypeInfo
+from mypy.nodes import MemberExpr, NameExpr, TypeInfo, Var
 from mypy.plugin import AttributeContext
 from mypy.typeanal import make_optional_type
 from mypy.types import AnyType, Instance, TupleType, TypeOfAny, get_proper_type
@@ -27,6 +27,12 @@ def transform_into_proper_attr_type(ctx: AttributeContext) -> MypyType:
 
     if isinstance(expr, NameExpr) and isinstance(expr.node, TypeInfo):
         node = expr.node
+    elif (
+        isinstance(expr, NameExpr)
+        and isinstance(expr.node, Var)
+        and isinstance(var_node_type := get_proper_type(expr.node.type), Instance)
+    ):
+        node = var_node_type.type
     else:
         ctx.api.fail("Unable to resolve type of choices property", ctx.context)
         return AnyType(TypeOfAny.from_error)

tests/assert_type/db/models/test_enums.py

@@ -76,6 +76,7 @@ class VoidChoices(BaseEmptyChoices):
 assert_type(Suit.CLUB.label, _StrOrPromise)
 assert_type(Suit.CLUB.value, int)
 assert_type(Suit.CLUB.do_not_call_in_templates, Literal[True])
+assert_type([e.value for e in Suit], list[int])
 
 assert_type(YearInSchool.names, list[str])
 assert_type(YearInSchool.labels, list[_StrOrPromise])
@@ -97,6 +98,7 @@ class VoidChoices(BaseEmptyChoices):
 assert_type(Vehicle.CAR.value, int)
 assert_type(Vehicle.CAR.do_not_call_in_templates, Literal[True])
 assert_type(Vehicle.__empty__, _StrOrPromise)
+assert_type([e.value for e in Vehicle], list[int])
 
 assert_type(Gender.names, list[str])
 assert_type(Gender.labels, list[_StrOrPromise])