Specific return types for values and values list (#53)

* Instead of using Literal types, overload QuerySet.values_list in the plugin. Fixes #43.

- Add a couple of extra type checks that Django makes:
  1) 'flat' and 'named' can't be used together.
  2) 'flat' is not valid when values_list is called with more than one field.

* Determine better row types for values_list/values based on fields specified.

- In the case of values_list, we use a Row type with either a single primitive, Tuple, or NamedTuple.
- In the case of values, we use a TypedDict.
- In both cases, Any is used as a fallback for individual fields if those fields cannot be resolved.

A couple other fixes I made along the way:
- Don't create reverse relation for ForeignKeys with related_name='+'
- Don't skip creating other related managers in AddRelatedManagers if a dynamic value is encountered
  for related_name parameter, or if the type cannot be determined.

* Fix for TypedDict so that they are considered anonymous.

* Clean up some comments.

* Implement making TypedDict anonymous in a way that doesn't crash sometimes.

* Fix flake8 errors.

* Remove even uglier hack about making TypedDict anonymous.

* Address review comments. Write a few better comments inside tests.

* Fix crash when running with mypyc ("interpreted classes cannot inherit from compiled") due to the way I extended TypedDictType.

- Implemented the hack in another way that works on mypyc.
- Added a couple extra tests of accessing 'id' / 'pk' via values_list.

* Fix flake8 errors.

* Support annotation expressions (use type Any) for TypedDicts row types returned by values_list.

- Bonus points: handle values_list gracefully (use type Any) where Tuples are returned
  where some of the fields arguments are not string literals.

Author: syastrov

django-stubs/__init__.pyi

@@ -1,7 +1,12 @@
-from typing import Any
+from typing import Any, NamedTuple
 from .utils.version import get_version as get_version
 
 VERSION: Any
 __version__: str
 
 def setup(set_prefix: bool = ...) -> None: ...
+
+# Used by mypy_django_plugin when returning a QuerySet row that is a NamedTuple where the field names are unknown
+class _NamedTupleAnyAttr(NamedTuple):
+    def __getattr__(self, item: str) -> Any: ...
+    def __setattr__(self, item: str, value: Any) -> None: ...

django-stubs/db/models/query.pyi

@@ -97,8 +97,9 @@ class QuerySet(Generic[_T, _Row], Collection[_Row], Sized):
     def raw(
         self, raw_query: str, params: Any = ..., translations: Optional[Dict[str, str]] = ..., using: None = ...
     ) -> RawQuerySet: ...
+    # The type of values may be overridden to be more specific in the mypy plugin, depending on the fields param
     def values(self, *fields: Union[str, Combinable], **expressions: Any) -> QuerySet[_T, Dict[str, Any]]: ...
-    # The type of values_list is overridden to be more specific in the mypy django plugin
+    # The type of values_list may be overridden to be more specific in the mypy plugin, depending on the fields param
     def values_list(
         self, *fields: Union[str, Combinable], flat: bool = ..., named: bool = ...
     ) -> QuerySet[_T, Any]: ...

mypy_django_plugin/helpers.py

@@ -1,14 +1,16 @@
 import typing
-from typing import Dict, Optional
+from collections import OrderedDict
+from typing import Dict, Optional, cast
 
-from mypy.checker import TypeChecker
+from mypy.checker import TypeChecker, gen_unique_name
+from mypy.mro import calculate_mro
 from mypy.nodes import (
     AssignmentStmt, ClassDef, Expression, ImportedName, Lvalue, MypyFile, NameExpr, SymbolNode, TypeInfo,
-)
+    SymbolTable, SymbolTableNode, Block, GDEF, MDEF, Var)
 from mypy.plugin import FunctionContext, MethodContext
 from mypy.types import (
     AnyType, Instance, NoneTyp, Type, TypeOfAny, TypeVarType, UnionType,
-)
+    TupleType, TypedDictType)
 
 MODEL_CLASS_FULLNAME = 'django.db.models.base.Model'
 FIELD_FULLNAME = 'django.db.models.fields.Field'
@@ -211,7 +213,7 @@ def extract_field_setter_type(tp: Instance) -> Optional[Type]:
     return None
 
 
-def extract_field_getter_type(tp: Instance) -> Optional[Type]:
+def extract_field_getter_type(tp: Type) -> Optional[Type]:
     if not isinstance(tp, Instance):
         return None
     if tp.type.has_base(FIELD_FULLNAME):
@@ -235,6 +237,10 @@ def get_fields_metadata(model: TypeInfo) -> Dict[str, typing.Any]:
     return get_django_metadata(model).setdefault('fields', {})
 
 
+def get_lookups_metadata(model: TypeInfo) -> Dict[str, typing.Any]:
+    return get_django_metadata(model).setdefault('lookups', {})
+
+
 def extract_explicit_set_type_of_model_primary_key(model: TypeInfo) -> Optional[Type]:
     """
     If field with primary_key=True is set on the model, extract its __set__ type.
@@ -296,3 +302,84 @@ def get_assigned_value_for_class(type_info: TypeInfo, name: str) -> Optional[Exp
         if isinstance(lvalue, NameExpr) and lvalue.name == name:
             return rvalue
     return None
+
+
+def is_field_nullable(model: TypeInfo, field_name: str) -> bool:
+    return get_fields_metadata(model).get(field_name, {}).get('null', False)
+
+
+def is_foreign_key(t: Type) -> bool:
+    if not isinstance(t, Instance):
+        return False
+    return has_any_of_bases(t.type, (FOREIGN_KEY_FULLNAME, ONETOONE_FIELD_FULLNAME))
+
+
+def build_class_with_annotated_fields(api: TypeChecker, base: Type, fields: 'OrderedDict[str, Type]',
+                                      name: str) -> Instance:
+    """Build an Instance with `name` that contains the specified `fields` as attributes and extends `base`."""
+    # Credit: This code is largely copied/modified from TypeChecker.intersect_instance_callable and
+    # NamedTupleAnalyzer.build_namedtuple_typeinfo
+
+    cur_module = cast(MypyFile, api.scope.stack[0])
+    gen_name = gen_unique_name(name, cur_module.names)
+
+    cdef = ClassDef(name, Block([]))
+    cdef.fullname = cur_module.fullname() + '.' + gen_name
+    info = TypeInfo(SymbolTable(), cdef, cur_module.fullname())
+    cdef.info = info
+    info.bases = [base]
+
+    def add_field(var: Var, is_initialized_in_class: bool = False,
+                  is_property: bool = False) -> None:
+        var.info = info
+        var.is_initialized_in_class = is_initialized_in_class
+        var.is_property = is_property
+        var._fullname = '%s.%s' % (info.fullname(), var.name())
+        info.names[var.name()] = SymbolTableNode(MDEF, var)
+
+    vars = [Var(item, typ) for item, typ in fields.items()]
+    for var in vars:
+        add_field(var, is_property=True)
+
+    calculate_mro(info)
+    info.calculate_metaclass_type()
+
+    cur_module.names[gen_name] = SymbolTableNode(GDEF, info, plugin_generated=True)
+    return Instance(info, [])
+
+
+def make_named_tuple(api: TypeChecker, fields: 'OrderedDict[str, Type]', name: str) -> Type:
+    if not fields:
+        # No fields specified, so fallback to a subclass of NamedTuple that allows
+        # __getattr__ / __setattr__ for any attribute name.
+        fallback = api.named_generic_type('django._NamedTupleAnyAttr', [])
+    else:
+        fallback = build_class_with_annotated_fields(
+            api=api,
+            base=api.named_generic_type('typing.NamedTuple', []),
+            fields=fields,
+            name=name
+        )
+    return TupleType(list(fields.values()), fallback=fallback)
+
+
+def make_typeddict(api: TypeChecker, fields: 'OrderedDict[str, Type]', required_keys: typing.Set[str]) -> Type:
+    object_type = api.named_generic_type('mypy_extensions._TypedDict', [])
+    typed_dict_type = TypedDictType(fields, required_keys=required_keys, fallback=object_type)
+    return typed_dict_type
+
+
+def make_tuple(api: TypeChecker, fields: typing.List[Type]) -> Type:
+    implicit_any = AnyType(TypeOfAny.special_form)
+    fallback = api.named_generic_type('builtins.tuple', [implicit_any])
+    return TupleType(fields, fallback=fallback)
+
+
+def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is_nullable: bool) -> Type:
+    node = type_info.get(private_field_name).node
+    if isinstance(node, Var):
+        descriptor_type = node.type
+        if is_nullable:
+            descriptor_type = make_optional(descriptor_type)
+        return descriptor_type
+    return AnyType(TypeOfAny.unannotated)

mypy_django_plugin/lookups.py

@@ -0,0 +1,150 @@
+import dataclasses
+from typing import Union, List
+
+from mypy.nodes import TypeInfo
+from mypy.plugin import CheckerPluginInterface
+from mypy.types import Type, Instance
+
+from mypy_django_plugin import helpers
+
+
+@dataclasses.dataclass
+class RelatedModelNode:
+    typ: Instance
+    is_nullable: bool
+
+
+@dataclasses.dataclass
+class FieldNode:
+    typ: Type
+
+
+LookupNode = Union[RelatedModelNode, FieldNode]
+
+
+class LookupException(Exception):
+    pass
+
+
+def resolve_lookup(api: CheckerPluginInterface, model_type_info: TypeInfo, lookup: str) -> List[LookupNode]:
+    """Resolve a lookup str to a list of LookupNodes.
+
+    Each node represents a part of the lookup (separated by "__"), in order.
+    Each node is the Model or Field that was resolved.
+
+    Raises LookupException if there were any issues resolving the lookup.
+    """
+    lookup_parts = lookup.split("__")
+
+    nodes = []
+    while lookup_parts:
+        lookup_part = lookup_parts.pop(0)
+
+        if not nodes:
+            current_node = None
+        else:
+            current_node = nodes[-1]
+
+        if current_node is None:
+            new_node = resolve_model_lookup(api, model_type_info, lookup_part)
+        elif isinstance(current_node, RelatedModelNode):
+            new_node = resolve_model_lookup(api, current_node.typ.type, lookup_part)
+        elif isinstance(current_node, FieldNode):
+            raise LookupException(f"Field lookups not yet supported for lookup {lookup}")
+        else:
+            raise LookupException(f"Unsupported node type: {type(current_node)}")
+        nodes.append(new_node)
+    return nodes
+
+
+def resolve_model_lookup(api: CheckerPluginInterface, model_type_info: TypeInfo,
+                         lookup: str) -> LookupNode:
+    """Resolve a lookup on the given model."""
+    if lookup == 'pk':
+        # Primary keys are special-cased
+        primary_key_type = helpers.extract_primary_key_type_for_get(model_type_info)
+        if primary_key_type:
+            return FieldNode(primary_key_type)
+        else:
+            # No PK, use the get type for AutoField as PK type.
+            autofield_info = api.lookup_typeinfo('django.db.models.fields.AutoField')
+            pk_type = helpers.get_private_descriptor_type(autofield_info, '_pyi_private_get_type',
+                                                          is_nullable=False)
+            return FieldNode(pk_type)
+
+    field_name = get_actual_field_name_for_lookup_field(lookup, model_type_info)
+
+    field_node = model_type_info.get(field_name)
+    if not field_node:
+        raise LookupException(
+            f'When resolving lookup "{lookup}", field "{field_name}" was not found in model {model_type_info.name()}')
+
+    if field_name.endswith('_id'):
+        field_name_without_id = field_name.rstrip('_id')
+        foreign_key_field = model_type_info.get(field_name_without_id)
+        if foreign_key_field is not None and helpers.is_foreign_key(foreign_key_field.type):
+            # Hack: If field ends with '_id' and there is a model field without the '_id' suffix, then use that field.
+            field_node = foreign_key_field
+            field_name = field_name_without_id
+
+    field_node_type = field_node.type
+    if field_node_type is None or not isinstance(field_node_type, Instance):
+        raise LookupException(
+            f'When resolving lookup "{lookup}", could not determine type for {model_type_info.name()}.{field_name}')
+
+    if helpers.is_foreign_key(field_node_type):
+        field_type = helpers.extract_field_getter_type(field_node_type)
+        is_nullable = helpers.is_optional(field_type)
+        if is_nullable:
+            field_type = helpers.make_required(field_type)
+
+        if isinstance(field_type, Instance):
+            return RelatedModelNode(typ=field_type, is_nullable=is_nullable)
+        else:
+            raise LookupException(f"Not an instance for field {field_type} lookup {lookup}")
+
+    field_type = helpers.extract_field_getter_type(field_node_type)
+
+    if field_type:
+        return FieldNode(typ=field_type)
+    else:
+        # Not a Field
+        if field_name == 'id':
+            # If no 'id' field was fouond, use an int
+            return FieldNode(api.named_generic_type('builtins.int', []))
+
+        related_manager_arg = None
+        if field_node_type.type.has_base(helpers.RELATED_MANAGER_CLASS_FULLNAME):
+            related_manager_arg = field_node_type.args[0]
+
+        if related_manager_arg is not None:
+            # Reverse relation
+            return RelatedModelNode(typ=related_manager_arg, is_nullable=True)
+        raise LookupException(
+            f'When resolving lookup "{lookup}", could not determine type for {model_type_info.name()}.{field_name}')
+
+
+def get_actual_field_name_for_lookup_field(lookup: str, model_type_info: TypeInfo) -> str:
+    """Attempt to find out the real field name if this lookup is a related_query_name (for reverse relations).
+
+    If it's not, return the original lookup.
+    """
+    lookups_metadata = helpers.get_lookups_metadata(model_type_info)
+    lookup_metadata = lookups_metadata.get(lookup)
+    if lookup_metadata is None:
+        # If not found on current model, look in all bases for their lookup metadata
+        for base in model_type_info.mro:
+            lookups_metadata = helpers.get_lookups_metadata(base)
+            lookup_metadata = lookups_metadata.get(lookup)
+            if lookup_metadata:
+                break
+    if not lookup_metadata:
+        lookup_metadata = {}
+    related_name = lookup_metadata.get('related_query_name_target', None)
+    if related_name:
+        # If the lookup is a related lookup, then look at the field specified by related_name.
+        # This is to support if related_query_name is set and differs from.
+        field_name = related_name
+    else:
+        field_name = lookup
+    return field_name

mypy_django_plugin/main.py

@@ -1,6 +1,5 @@
-from functools import partial
-
 import os
+from functools import partial
 from typing import Callable, Dict, Optional, Union, cast
 
 from mypy.checker import TypeChecker
@@ -23,6 +22,7 @@
     determine_model_cls_from_string_for_migrations, get_string_value_from_expr,
 )
 from mypy_django_plugin.transformers.models import process_model_class
+from mypy_django_plugin.transformers.queryset import extract_proper_type_for_values_and_values_list
 from mypy_django_plugin.transformers.settings import (
     AddSettingValuesToDjangoConfObject, get_settings_metadata,
 )
@@ -165,7 +165,7 @@ def extract_and_return_primary_key_of_bound_related_field_parameter(ctx: Attribu
         if primary_key_type:
             return primary_key_type
 
-    is_nullable = helpers.get_fields_metadata(ctx.type.type).get(field_name, {}).get('null', False)
+    is_nullable = helpers.is_field_nullable(ctx.type.type, field_name)
     if is_nullable:
         return helpers.make_optional(ctx.default_attr_type)
 
@@ -292,7 +292,10 @@ def __init__(self, options: Options) -> None:
         if self.django_settings_module:
             settings_modules.append(self.django_settings_module)
 
-        monkeypatch.add_modules_as_a_source_seed_files(settings_modules)
+        auto_imports = ['mypy_extensions']
+        auto_imports.extend(settings_modules)
+
+        monkeypatch.add_modules_as_a_source_seed_files(auto_imports)
         monkeypatch.inject_modules_as_dependencies_for_django_conf_settings(settings_modules)
 
     def _get_current_model_bases(self) -> Dict[str, int]:
@@ -359,10 +362,10 @@ def get_method_hook(self, fullname: str
             if sym and isinstance(sym.node, TypeInfo) and sym.node.has_base(helpers.FORM_MIXIN_CLASS_FULLNAME):
                 return extract_proper_type_for_get_form
 
-        if method_name == 'values_list':
+        if method_name in ('values', 'values_list'):
             sym = self.lookup_fully_qualified(class_name)
             if sym and isinstance(sym.node, TypeInfo) and sym.node.has_base(helpers.QUERYSET_CLASS_FULLNAME):
-                return extract_proper_type_for_values_list
+                return partial(extract_proper_type_for_values_and_values_list, method_name)
 
         if fullname in {'django.apps.registry.Apps.get_model',
                         'django.db.migrations.state.StateApps.get_model'}:

mypy_django_plugin/transformers/fields.py

@@ -1,10 +1,10 @@
 from typing import Optional, cast
 
 from mypy.checker import TypeChecker
-from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo, Var
+from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo
 from mypy.plugin import FunctionContext
 from mypy.types import (
-    AnyType, CallableType, Instance, TupleType, Type, TypeOfAny, UnionType,
+    AnyType, CallableType, Instance, TupleType, Type, UnionType,
 )
 
 from mypy_django_plugin import helpers
@@ -88,23 +88,13 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext) -> Type:
     return helpers.reparametrize_instance(ctx.default_return_type, new_args=args)
 
 
-def get_private_descriptor_type(type_info: TypeInfo, private_field_name: str, is_nullable: bool) -> Type:
-    node = type_info.get(private_field_name).node
-    if isinstance(node, Var):
-        descriptor_type = node.type
-        if is_nullable:
-            descriptor_type = helpers.make_optional(descriptor_type)
-        return descriptor_type
-    return AnyType(TypeOfAny.unannotated)
-
-
 def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance:
     default_return_type = cast(Instance, ctx.default_return_type)
     is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'null'))
-    set_type = get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type',
-                                           is_nullable=is_nullable)
-    get_type = get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type',
-                                           is_nullable=is_nullable)
+    set_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type',
+                                                   is_nullable=is_nullable)
+    get_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type',
+                                                   is_nullable=is_nullable)
     return helpers.reparametrize_instance(default_return_type, [set_type, get_type])