Add support for BaseManager.from_queryset() (#251)

* add support for BaseManager.from_queryset()

* cleanups

* lint fixes

Author: mkurnikov

mypy_django_plugin/lib/helpers.py

@@ -1,6 +1,6 @@
 from collections import OrderedDict
 from typing import (
-    TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Union, cast,
+    TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union,
 )
 
 from django.db.models.fields import Field
@@ -10,13 +10,15 @@
 from mypy.checker import TypeChecker
 from mypy.mro import calculate_mro
 from mypy.nodes import (
-    GDEF, MDEF, Block, ClassDef, Expression, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode, SymbolTable,
-    SymbolTableNode, TypeInfo, Var,
+    GDEF, MDEF, Argument, Block, ClassDef, Expression, FuncDef, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolNode,
+    SymbolTable, SymbolTableNode, TypeInfo, Var,
 )
 from mypy.plugin import (
-    AttributeContext, CheckerPluginInterface, FunctionContext, MethodContext,
+    AttributeContext, CheckerPluginInterface, ClassDefContext, DynamicClassDefContext, FunctionContext, MethodContext,
 )
-from mypy.types import AnyType, Instance, NoneTyp, TupleType
+from mypy.plugins.common import add_method
+from mypy.semanal import SemanticAnalyzer
+from mypy.types import AnyType, CallableType, Instance, NoneTyp, TupleType
 from mypy.types import Type as MypyType
 from mypy.types import TypedDictType, TypeOfAny, UnionType
 
@@ -55,7 +57,7 @@ def lookup_fully_qualified_generic(name: str, all_modules: Dict[str, MypyFile])
     return sym.node
 
 
-def lookup_fully_qualified_typeinfo(api: TypeChecker, fullname: str) -> Optional[TypeInfo]:
+def lookup_fully_qualified_typeinfo(api: Union[TypeChecker, SemanticAnalyzer], fullname: str) -> Optional[TypeInfo]:
     node = lookup_fully_qualified_generic(fullname, api.modules)
     if not isinstance(node, TypeInfo):
         return None
@@ -173,8 +175,11 @@ def get_nested_meta_node_for_current_class(info: TypeInfo) -> Optional[TypeInfo]
     return None
 
 
-def add_new_class_for_module(module: MypyFile, name: str, bases: List[Instance],
-                             fields: 'OrderedDict[str, MypyType]') -> TypeInfo:
+def add_new_class_for_module(module: MypyFile,
+                             name: str,
+                             bases: List[Instance],
+                             fields: Optional[Dict[str, MypyType]] = None
+                             ) -> TypeInfo:
     new_class_unique_name = checker.gen_unique_name(name, module.names)
 
     # make new class expression
@@ -188,11 +193,12 @@ def add_new_class_for_module(module: MypyFile, name: str, bases: List[Instance],
     new_typeinfo.calculate_metaclass_type()
 
     # add fields
-    for field_name, field_type in fields.items():
-        var = Var(field_name, type=field_type)
-        var.info = new_typeinfo
-        var._fullname = new_typeinfo.fullname + '.' + field_name
-        new_typeinfo.names[field_name] = SymbolTableNode(MDEF, var, plugin_generated=True)
+    if fields:
+        for field_name, field_type in fields.items():
+            var = Var(field_name, type=field_type)
+            var.info = new_typeinfo
+            var._fullname = new_typeinfo.fullname + '.' + field_name
+            new_typeinfo.names[field_name] = SymbolTableNode(MDEF, var, plugin_generated=True)
 
     classdef.info = new_typeinfo
     module.names[new_class_unique_name] = SymbolTableNode(GDEF, new_typeinfo, plugin_generated=True)
@@ -269,10 +275,16 @@ def resolve_string_attribute_value(attr_expr: Expression, ctx: Union[FunctionCon
     return None
 
 
+def get_semanal_api(ctx: Union[ClassDefContext, DynamicClassDefContext]) -> SemanticAnalyzer:
+    if not isinstance(ctx.api, SemanticAnalyzer):
+        raise ValueError('Not a SemanticAnalyzer')
+    return ctx.api
+
+
 def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionContext]) -> TypeChecker:
     if not isinstance(ctx.api, TypeChecker):
         raise ValueError('Not a TypeChecker')
-    return cast(TypeChecker, ctx.api)
+    return ctx.api
 
 
 def is_model_subclass_info(info: TypeInfo, django_context: 'DjangoContext') -> bool:
@@ -298,3 +310,28 @@ def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType) -> No
     var.is_inferred = True
     info.names[name] = SymbolTableNode(MDEF, var,
                                        plugin_generated=True)
+
+
+def _prepare_new_method_arguments(node: FuncDef) -> Tuple[List[Argument], MypyType]:
+    arguments = []
+    for argument in node.arguments[1:]:
+        if argument.type_annotation is None:
+            argument.type_annotation = AnyType(TypeOfAny.unannotated)
+        arguments.append(argument)
+
+    if isinstance(node.type, CallableType):
+        return_type = node.type.ret_type
+    else:
+        return_type = AnyType(TypeOfAny.unannotated)
+
+    return arguments, return_type
+
+
+def copy_method_to_another_class(ctx: ClassDefContext, self_type: Instance,
+                                 new_method_name: str, method_node: FuncDef) -> None:
+    arguments, return_type = _prepare_new_method_arguments(method_node)
+    add_method(ctx,
+               new_method_name,
+               args=arguments,
+               return_type=return_type,
+               self_type=self_type)

mypy_django_plugin/main.py

@@ -7,7 +7,7 @@
 from mypy.nodes import MypyFile, TypeInfo
 from mypy.options import Options
 from mypy.plugin import (
-    AttributeContext, ClassDefContext, FunctionContext, MethodContext, Plugin,
+    AttributeContext, ClassDefContext, DynamicClassDefContext, FunctionContext, MethodContext, Plugin,
 )
 from mypy.types import Type as MypyType
 
@@ -17,6 +17,9 @@
 from mypy_django_plugin.transformers import (
     fields, forms, init_create, meta, querysets, request, settings,
 )
+from mypy_django_plugin.transformers.managers import (
+    create_new_manager_class_from_from_queryset_method,
+)
 from mypy_django_plugin.transformers.models import process_model_class
 
 
@@ -242,6 +245,15 @@ def get_attribute_hook(self, fullname: str
             return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context)
         return None
 
+    def get_dynamic_class_hook(self, fullname: str
+                               ) -> Optional[Callable[[DynamicClassDefContext], None]]:
+        if fullname.endswith('from_queryset'):
+            class_name, _, _ = fullname.rpartition('.')
+            info = self._get_typeinfo_or_none(class_name)
+            if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
+                return create_new_manager_class_from_from_queryset_method
+        return None
+
 
 def plugin(version):
     return NewSemanalDjangoPlugin

mypy_django_plugin/transformers/managers.py

@@ -0,0 +1,71 @@
+from mypy.nodes import (
+    GDEF, FuncDef, MemberExpr, NameExpr, StrExpr, SymbolTableNode, TypeInfo,
+)
+from mypy.plugin import ClassDefContext, DynamicClassDefContext
+from mypy.types import AnyType, Instance, TypeOfAny
+
+from mypy_django_plugin.lib import helpers
+
+
+def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None:
+    semanal_api = helpers.get_semanal_api(ctx)
+
+    assert isinstance(ctx.call.callee, MemberExpr)
+    assert isinstance(ctx.call.callee.expr, NameExpr)
+    base_manager_info = ctx.call.callee.expr.node
+    if base_manager_info is None:
+        if not semanal_api.final_iteration:
+            semanal_api.defer()
+        return
+
+    assert isinstance(base_manager_info, TypeInfo)
+    new_manager_info = semanal_api.basic_new_typeinfo(ctx.name,
+                                                      basetype_or_fallback=Instance(base_manager_info,
+                                                                                    [AnyType(TypeOfAny.unannotated)]))
+    new_manager_info.line = ctx.call.line
+    new_manager_info.defn.line = ctx.call.line
+    new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type()
+
+    current_module = semanal_api.cur_mod_node
+    current_module.names[ctx.name] = SymbolTableNode(GDEF, new_manager_info,
+                                                     plugin_generated=True)
+    passed_queryset = ctx.call.args[0]
+    assert isinstance(passed_queryset, NameExpr)
+
+    derived_queryset_fullname = passed_queryset.fullname
+    assert derived_queryset_fullname is not None
+
+    sym = semanal_api.lookup_fully_qualified_or_none(derived_queryset_fullname)
+    assert sym is not None
+    if sym.node is None:
+        if not semanal_api.final_iteration:
+            semanal_api.defer()
+        else:
+            # inherit from Any to prevent false-positives, if queryset class cannot be resolved
+            new_manager_info.fallback_to_any = True
+        return
+
+    derived_queryset_info = sym.node
+    assert isinstance(derived_queryset_info, TypeInfo)
+
+    if len(ctx.call.args) > 1:
+        expr = ctx.call.args[1]
+        assert isinstance(expr, StrExpr)
+        custom_manager_generated_name = expr.value
+    else:
+        custom_manager_generated_name = base_manager_info.name + 'From' + derived_queryset_info.name
+
+    custom_manager_generated_fullname = '.'.join(['django.db.models.manager', custom_manager_generated_name])
+    if 'from_queryset_managers' not in base_manager_info.metadata:
+        base_manager_info.metadata['from_queryset_managers'] = {}
+    base_manager_info.metadata['from_queryset_managers'][custom_manager_generated_fullname] = new_manager_info.fullname
+
+    class_def_context = ClassDefContext(cls=new_manager_info.defn,
+                                        reason=ctx.call, api=semanal_api)
+    self_type = Instance(new_manager_info, [])
+    for name, sym in derived_queryset_info.names.items():
+        if isinstance(sym.node, FuncDef):
+            helpers.copy_method_to_another_class(class_def_context,
+                                                 self_type,
+                                                 new_method_name=name,
+                                                 method_node=sym.node)