Fixes for Manager.from_queryset(), support for QuerySet.as_manager()

mypy_django_plugin/django/context.py

@@ -21,7 +21,7 @@
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, UnionType
 
-from mypy_django_plugin.lib import fullnames, helpers
+from mypy_django_plugin.lib import chk_helpers, fullnames, helpers
 
 try:
     from django.contrib.postgres.fields import ArrayField
@@ -356,11 +356,11 @@ def resolve_lookup_expected_type(self, ctx: MethodContext, model_cls: Type[Model
                 return AnyType(TypeOfAny.explicit)
 
         if lookup_cls is None or isinstance(lookup_cls, Exact):
-            return self.get_field_lookup_exact_type(helpers.get_typechecker_api(ctx), field)
+            return self.get_field_lookup_exact_type(chk_helpers.get_typechecker_api(ctx), field)
 
         assert lookup_cls is not None
 
-        lookup_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), lookup_cls)
+        lookup_info = helpers.lookup_class_typeinfo(chk_helpers.get_typechecker_api(ctx), lookup_cls)
         if lookup_info is None:
             return AnyType(TypeOfAny.explicit)
 
@@ -370,7 +370,7 @@ def resolve_lookup_expected_type(self, ctx: MethodContext, model_cls: Type[Model
                 # if it's Field, consider lookup_type a __get__ of current field
                 if (isinstance(lookup_type, Instance)
                         and lookup_type.type.fullname == fullnames.FIELD_FULLNAME):
-                    field_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), field.__class__)
+                    field_info = helpers.lookup_class_typeinfo(chk_helpers.get_typechecker_api(ctx), field.__class__)
                     if field_info is None:
                         return AnyType(TypeOfAny.explicit)
                     lookup_type = helpers.get_private_descriptor_type(field_info, '_pyi_private_get_type',

mypy_django_plugin/lib/chk_helpers.py

@@ -0,0 +1,120 @@
+from typing import Dict, List, Optional, Set, Union
+
+from mypy import checker
+from mypy.checker import TypeChecker
+from mypy.nodes import (
+    GDEF, MDEF, Expression, MypyFile, SymbolTableNode, TypeInfo, Var,
+)
+from mypy.plugin import (
+    AttributeContext, CheckerPluginInterface, FunctionContext, MethodContext,
+)
+from mypy.types import AnyType, Instance, TupleType
+from mypy.types import Type as MypyType
+from mypy.types import TypedDictType, TypeOfAny
+
+from mypy_django_plugin.lib import helpers
+
+
+def add_new_class_for_current_module(current_module: MypyFile,
+                                     name: str,
+                                     bases: List[Instance],
+                                     fields: Optional[Dict[str, MypyType]] = None
+                                     ) -> TypeInfo:
+    new_class_unique_name = checker.gen_unique_name(name, current_module.names)
+    new_typeinfo = helpers.new_typeinfo(new_class_unique_name,
+                                        bases=bases,
+                                        module_name=current_module.fullname)
+    # new_typeinfo = helpers.make_new_typeinfo_in_current_module(new_class_unique_name,
+    #                                                            bases=bases,
+    #                                                            current_module_fullname=current_module.fullname)
+    # add fields
+    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)
+
+    current_module.names[new_class_unique_name] = SymbolTableNode(GDEF, new_typeinfo, plugin_generated=True)
+    current_module.defs.append(new_typeinfo.defn)
+    return new_typeinfo
+
+
+def make_oneoff_named_tuple(api: TypeChecker, name: str, fields: 'Dict[str, MypyType]') -> TupleType:
+    current_module = helpers.get_current_module(api)
+    namedtuple_info = add_new_class_for_current_module(current_module, name,
+                                                       bases=[api.named_generic_type('typing.NamedTuple', [])],
+                                                       fields=fields)
+    return TupleType(list(fields.values()), fallback=Instance(namedtuple_info, []))
+
+
+def make_tuple(api: 'TypeChecker', fields: List[MypyType]) -> TupleType:
+    # fallback for tuples is any builtins.tuple instance
+    fallback = api.named_generic_type('builtins.tuple',
+                                      [AnyType(TypeOfAny.special_form)])
+    return TupleType(fields, fallback=fallback)
+
+
+def make_oneoff_typeddict(api: CheckerPluginInterface, fields: 'Dict[str, MypyType]',
+                          required_keys: Set[str]) -> TypedDictType:
+    object_type = api.named_generic_type('mypy_extensions._TypedDict', [])
+    typed_dict_type = TypedDictType(fields,  # type: ignore
+                                    required_keys=required_keys,
+                                    fallback=object_type)
+    return typed_dict_type
+
+
+def get_typechecker_api(ctx: Union[AttributeContext, MethodContext, FunctionContext]) -> TypeChecker:
+    if not isinstance(ctx.api, TypeChecker):
+        raise ValueError('Not a TypeChecker')
+    return ctx.api
+
+
+def check_types_compatible(ctx: Union[FunctionContext, MethodContext],
+                           *, expected_type: MypyType, actual_type: MypyType, error_message: str) -> None:
+    api = get_typechecker_api(ctx)
+    api.check_subtype(actual_type, expected_type,
+                      ctx.context, error_message,
+                      'got', 'expected')
+
+
+def get_call_argument_by_name(ctx: Union[FunctionContext, MethodContext], name: str) -> Optional[Expression]:
+    """
+    Return the expression for the specific argument.
+    This helper should only be used with non-star arguments.
+    """
+    if name not in ctx.callee_arg_names:
+        return None
+    idx = ctx.callee_arg_names.index(name)
+    args = ctx.args[idx]
+    if len(args) != 1:
+        # Either an error or no value passed.
+        return None
+    return args[0]
+
+
+def get_call_argument_type_by_name(ctx: Union[FunctionContext, MethodContext], name: str) -> Optional[MypyType]:
+    """Return the type for the specific argument.
+
+    This helper should only be used with non-star arguments.
+    """
+    if name not in ctx.callee_arg_names:
+        return None
+    idx = ctx.callee_arg_names.index(name)
+    arg_types = ctx.arg_types[idx]
+    if len(arg_types) != 1:
+        # Either an error or no value passed.
+        return None
+    return arg_types[0]
+
+
+def add_new_sym_for_info(info: TypeInfo, *, name: str, sym_type: MypyType) -> None:
+    # type=: type of the variable itself
+    var = Var(name=name, type=sym_type)
+    # var.info: type of the object variable is bound to
+    var.info = info
+    var._fullname = info.fullname + '.' + name
+    var.is_initialized_in_class = True
+    var.is_inferred = True
+    info.names[name] = SymbolTableNode(MDEF, var,
+                                       plugin_generated=True)