move querysets to new API

Author: mkurnikov

mypy_django_plugin/main.py

@@ -1,5 +1,4 @@
 import configparser
-from functools import partial
 from typing import Callable, Dict, List, Optional, Tuple
 
 from django.db.models.fields.related import RelatedField
@@ -13,7 +12,6 @@
 
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
-from mypy_django_plugin.transformers import querysets
 from mypy_django_plugin.transformers2.fields import FieldContructorCallback
 from mypy_django_plugin.transformers2.forms import (
     FormCallback, GetFormCallback, GetFormClassCallback,
@@ -26,6 +24,9 @@
 from mypy_django_plugin.transformers2.orm_lookups import (
     QuerySetFilterTypecheckCallback,
 )
+from mypy_django_plugin.transformers2.querysets import (
+    QuerySetValuesCallback, QuerySetValuesListCallback,
+)
 from mypy_django_plugin.transformers2.related_managers import (
     GetRelatedManagerCallback,
 )
@@ -176,7 +177,6 @@ def get_function_hook(self, fullname: str
 
             if self.django_context.is_model_subclass(info):
                 return ModelInitCallback(self)
-                # return partial(init_create.redefine_and_typecheck_model_init, django_context=self.django_context)
         return None
 
     def get_method_hook(self, fullname: str
@@ -195,12 +195,14 @@ def get_method_hook(self, fullname: str
         if method_name == 'values':
             info = self._get_typeinfo_or_none(class_fullname)
             if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
-                return partial(querysets.extract_proper_type_queryset_values, django_context=self.django_context)
+                return QuerySetValuesCallback(self)
+                # return partial(querysets.extract_proper_type_queryset_values, django_context=self.django_context)
 
         if method_name == 'values_list':
             info = self._get_typeinfo_or_none(class_fullname)
             if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
-                return partial(querysets.extract_proper_type_queryset_values_list, django_context=self.django_context)
+                return QuerySetValuesListCallback(self)
+                # return partial(querysets.extract_proper_type_queryset_values_list, django_context=self.django_context)
 
         if method_name == 'get_field':
             info = self._get_typeinfo_or_none(class_fullname)

mypy_django_plugin/transformers/__init__.py

@@ -1,12 +1,12 @@
-from collections import OrderedDict
+import collections
 from typing import List, Optional, Sequence, Type
 
 from django.core.exceptions import FieldError
 from django.db.models.base import Model
 from django.db.models.fields.related import RelatedField
 from django.db.models.fields.reverse_related import ForeignObjectRel
 from mypy.nodes import Expression, NameExpr
-from mypy.plugin import FunctionContext, MethodContext
+from mypy.plugin import MethodContext
 from mypy.types import AnyType, Instance
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny
@@ -17,27 +17,6 @@
 from mypy_django_plugin.lib import chk_helpers, fullnames, helpers
 
 
-def _extract_model_type_from_queryset(queryset_type: Instance) -> Optional[Instance]:
-    for base_type in [queryset_type, *queryset_type.type.bases]:
-        if (len(base_type.args)
-                and isinstance(base_type.args[0], Instance)
-                and base_type.args[0].type.has_base(fullnames.MODEL_CLASS_FULLNAME)):
-            return base_type.args[0]
-    return None
-
-
-def determine_proper_manager_type(ctx: FunctionContext) -> MypyType:
-    default_return_type = ctx.default_return_type
-    assert isinstance(default_return_type, Instance)
-
-    outer_model_info = chk_helpers.get_typechecker_api(ctx).scope.active_class()
-    if (outer_model_info is None
-            or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME)):
-        return default_return_type
-
-    return helpers.reparametrize_instance(default_return_type, [Instance(outer_model_info, [])])
-
-
 def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
                                *, method: str, lookup: str) -> Optional[MypyType]:
     try:
@@ -60,6 +39,16 @@ def get_field_type_from_lookup(ctx: MethodContext, django_context: DjangoContext
     return field_get_type
 
 
+def resolve_field_lookups(lookup_exprs: Sequence[Expression], django_context: DjangoContext) -> Optional[List[str]]:
+    field_lookups = []
+    for field_lookup_expr in lookup_exprs:
+        field_lookup = helpers.resolve_string_attribute_value(field_lookup_expr, django_context)
+        if field_lookup is None:
+            return None
+        field_lookups.append(field_lookup)
+    return field_lookups
+
+
 def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext, model_cls: Type[Model],
                              flat: bool, named: bool) -> MypyType:
     field_lookups = resolve_field_lookups(ctx.args[0], django_context)
@@ -75,7 +64,7 @@ def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext,
             assert lookup_type is not None
             return lookup_type
         elif named:
-            column_types: 'OrderedDict[str, MypyType]' = OrderedDict()
+            column_types = collections.OrderedDict()
             for field in django_context.get_model_fields(model_cls):
                 column_type = django_context.get_field_get_type(typechecker_api, field,
                                                                 method='values_list')
@@ -91,7 +80,7 @@ def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext,
         typechecker_api.fail("'flat' is not valid when 'values_list' is called with more than one field", ctx.context)
         return AnyType(TypeOfAny.from_error)
 
-    column_types = OrderedDict()
+    column_types = collections.OrderedDict()
     for field_lookup in field_lookups:
         lookup_field_type = get_field_type_from_lookup(ctx, django_context, model_cls,
                                                        lookup=field_lookup, method='values_list')
@@ -110,83 +99,83 @@ def get_values_list_row_type(ctx: MethodContext, django_context: DjangoContext,
     return row_type
 
 
-def extract_proper_type_queryset_values_list(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
-    # called on the Instance, returns QuerySet of something
-    assert isinstance(ctx.type, Instance)
-    assert isinstance(ctx.default_return_type, Instance)
+class QuerySetMethodCallback(helpers.GetMethodCallback):
+    def current_model_type(self) -> Optional[Instance]:
+        for base_type in [self.callee_type, *self.callee_type.type.bases]:
+            if (len(base_type.args)
+                    and isinstance(base_type.args[0], Instance)
+                    and base_type.args[0].type.has_base(fullnames.MODEL_CLASS_FULLNAME)):
+                return base_type.args[0]
+        return None
 
-    model_type = _extract_model_type_from_queryset(ctx.type)
-    if model_type is None:
-        return AnyType(TypeOfAny.from_omitted_generics)
 
-    model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname)
-    if model_cls is None:
-        return ctx.default_return_type
+class QuerySetValuesCallback(QuerySetMethodCallback):
+    def get_method_return_type(self) -> MypyType:
+        assert isinstance(self.default_return_type, Instance)
 
-    flat_expr = chk_helpers.get_call_argument_by_name(ctx, 'flat')
-    if flat_expr is not None and isinstance(flat_expr, NameExpr):
-        flat = helpers.parse_bool(flat_expr)
-    else:
-        flat = False
+        model_type = self.current_model_type()
+        if model_type is None:
+            return AnyType(TypeOfAny.from_omitted_generics)
 
-    named_expr = chk_helpers.get_call_argument_by_name(ctx, 'named')
-    if named_expr is not None and isinstance(named_expr, NameExpr):
-        named = helpers.parse_bool(named_expr)
-    else:
-        named = False
+        model_cls = self.django_context.get_model_class_by_fullname(model_type.type.fullname)
+        if model_cls is None:
+            return self.default_return_type
 
-    if flat and named:
-        ctx.api.fail("'flat' and 'named' can't be used together", ctx.context)
-        return helpers.reparametrize_instance(ctx.default_return_type, [model_type, AnyType(TypeOfAny.from_error)])
+        field_lookups = resolve_field_lookups(self.ctx.args[0], self.django_context)
+        if field_lookups is None:
+            return AnyType(TypeOfAny.from_error)
 
-    # account for possible None
-    flat = flat or False
-    named = named or False
+        if len(field_lookups) == 0:
+            for field in self.django_context.get_model_fields(model_cls):
+                field_lookups.append(field.attname)
 
-    row_type = get_values_list_row_type(ctx, django_context, model_cls,
-                                        flat=flat, named=named)
-    return helpers.reparametrize_instance(ctx.default_return_type, [model_type, row_type])
+        column_types = collections.OrderedDict()
+        for field_lookup in field_lookups:
+            field_lookup_type = get_field_type_from_lookup(self.ctx, self.django_context, model_cls,
+                                                           lookup=field_lookup, method='values')
+            if field_lookup_type is None:
+                return helpers.reparametrize_instance(self.default_return_type,
+                                                      [model_type, AnyType(TypeOfAny.from_error)])
 
+            column_types[field_lookup] = field_lookup_type
 
-def resolve_field_lookups(lookup_exprs: Sequence[Expression], django_context: DjangoContext) -> Optional[List[str]]:
-    field_lookups = []
-    for field_lookup_expr in lookup_exprs:
-        field_lookup = helpers.resolve_string_attribute_value(field_lookup_expr, django_context)
-        if field_lookup is None:
-            return None
-        field_lookups.append(field_lookup)
-    return field_lookups
+        row_type = chk_helpers.make_oneoff_typeddict(self.ctx.api, column_types, set(column_types.keys()))
+        return helpers.reparametrize_instance(self.default_return_type, [model_type, row_type])
 
 
-def extract_proper_type_queryset_values(ctx: MethodContext, django_context: DjangoContext) -> MypyType:
-    # called on QuerySet, return QuerySet of something
-    assert isinstance(ctx.type, Instance)
-    assert isinstance(ctx.default_return_type, Instance)
+class QuerySetValuesListCallback(QuerySetMethodCallback):
+    def get_method_return_type(self) -> MypyType:
+        # called on the Instance, returns QuerySet of something
+        assert isinstance(self.default_return_type, Instance)
 
-    model_type = _extract_model_type_from_queryset(ctx.type)
-    if model_type is None:
-        return AnyType(TypeOfAny.from_omitted_generics)
+        model_type = self.current_model_type()
+        if model_type is None:
+            return AnyType(TypeOfAny.from_omitted_generics)
 
-    model_cls = django_context.get_model_class_by_fullname(model_type.type.fullname)
-    if model_cls is None:
-        return ctx.default_return_type
+        model_cls = self.django_context.get_model_class_by_fullname(model_type.type.fullname)
+        if model_cls is None:
+            return self.default_return_type
 
-    field_lookups = resolve_field_lookups(ctx.args[0], django_context)
-    if field_lookups is None:
-        return AnyType(TypeOfAny.from_error)
+        flat_expr = chk_helpers.get_call_argument_by_name(self.ctx, 'flat')
+        if flat_expr is not None and isinstance(flat_expr, NameExpr):
+            flat = helpers.parse_bool(flat_expr)
+        else:
+            flat = False
 
-    if len(field_lookups) == 0:
-        for field in django_context.get_model_fields(model_cls):
-            field_lookups.append(field.attname)
+        named_expr = chk_helpers.get_call_argument_by_name(self.ctx, 'named')
+        if named_expr is not None and isinstance(named_expr, NameExpr):
+            named = helpers.parse_bool(named_expr)
+        else:
+            named = False
 
-    column_types: 'OrderedDict[str, MypyType]' = OrderedDict()
-    for field_lookup in field_lookups:
-        field_lookup_type = get_field_type_from_lookup(ctx, django_context, model_cls,
-                                                       lookup=field_lookup, method='values')
-        if field_lookup_type is None:
-            return helpers.reparametrize_instance(ctx.default_return_type, [model_type, AnyType(TypeOfAny.from_error)])
+        if flat and named:
+            self.ctx.api.fail("'flat' and 'named' can't be used together", self.ctx.context)
+            return helpers.reparametrize_instance(self.default_return_type, [model_type, AnyType(TypeOfAny.from_error)])
 
-        column_types[field_lookup] = field_lookup_type
+        # account for possible None
+        flat = flat or False
+        named = named or False
 
-    row_type = chk_helpers.make_oneoff_typeddict(ctx.api, column_types, set(column_types.keys()))
-    return helpers.reparametrize_instance(ctx.default_return_type, [model_type, row_type])
+        row_type = get_values_list_row_type(self.ctx, self.django_context, model_cls,
+                                            flat=flat, named=named)
+        return helpers.reparametrize_instance(self.default_return_type, [model_type, row_type])