New semanal-progress rebased on master

django-stubs/db/models/fields/__init__.pyi

@@ -406,4 +406,5 @@ class DurationField(Field[_ST, _GT]):
     _pyi_private_get_type: timedelta
 
 class BigAutoField(AutoField[_ST, _GT]): ...
+class SmallAutoField(AutoField[_ST, _GT]): ...
 class CommaSeparatedIntegerField(CharField[_ST, _GT]): ...

mypy_django_plugin/django/context.py

@@ -16,12 +16,13 @@
 from django.db.models.sql.query import Query
 from django.utils.functional import cached_property
 from mypy.checker import TypeChecker
+from mypy.nodes import TypeInfo
 from mypy.plugin import MethodContext
 from mypy.types import AnyType, Instance
 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
@@ -52,14 +53,6 @@ def initialize_django(settings_module: str) -> Tuple['Apps', 'LazySettings']:
         # add current directory to sys.path
         sys.path.append(os.getcwd())
 
-        def noop_class_getitem(cls, key):
-            return cls
-
-        from django.db import models
-
-        models.QuerySet.__class_getitem__ = classmethod(noop_class_getitem)  # type: ignore
-        models.Manager.__class_getitem__ = classmethod(noop_class_getitem)  # type: ignore
-
         from django.conf import settings
         from django.apps import apps
 
@@ -119,10 +112,10 @@ def get_model_fields(self, model_cls: Type[Model]) -> Iterator[Field]:
             if isinstance(field, Field):
                 yield field
 
-    def get_model_relations(self, model_cls: Type[Model]) -> Iterator[ForeignObjectRel]:
-        for field in model_cls._meta.get_fields():
-            if isinstance(field, ForeignObjectRel):
-                yield field
+    def get_model_relations(self, model_cls: Type[Model]) -> Iterator[Tuple[Optional[str], ForeignObjectRel]]:
+        for relation in model_cls._meta.get_fields():
+            if isinstance(relation, ForeignObjectRel):
+                yield relation.get_accessor_name(), relation
 
     def get_field_lookup_exact_type(self, api: TypeChecker, field: Union[Field, ForeignObjectRel]) -> MypyType:
         if isinstance(field, (RelatedField, ForeignObjectRel)):
@@ -222,11 +215,15 @@ def all_registered_model_classes(self) -> Set[Type[models.Model]]:
     def all_registered_model_class_fullnames(self) -> Set[str]:
         return {helpers.get_class_fullname(cls) for cls in self.all_registered_model_classes}
 
+    def is_model_subclass(self, info: TypeInfo) -> bool:
+        return (info.fullname in self.all_registered_model_class_fullnames
+                or info.has_base(fullnames.MODEL_CLASS_FULLNAME))
+
     def get_attname(self, field: Field) -> str:
         attname = field.attname
         return attname
 
-    def get_field_nullability(self, field: Union[Field, ForeignObjectRel], method: Optional[str]) -> bool:
+    def get_field_nullability(self, field: Union[Field, ForeignObjectRel], method: Optional[str] = None) -> bool:
         nullable = field.null
         if not nullable and isinstance(field, CharField) and field.blank:
             return True
@@ -356,11 +353,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 +367,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,126 @@
+from typing import Dict, List, Optional, Set, Union
+
+from mypy import checker
+from mypy.checker import TypeChecker
+from mypy.mro import calculate_mro
+from mypy.nodes import (
+    GDEF, MDEF, Block, ClassDef, Expression, MypyFile, SymbolTable, 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)
+
+    # make new class expression
+    classdef = ClassDef(new_class_unique_name, Block([]))
+    classdef.fullname = current_module.fullname + '.' + new_class_unique_name
+
+    # make new TypeInfo
+    new_typeinfo = TypeInfo(SymbolTable(), classdef, current_module.fullname)
+    new_typeinfo.bases = bases
+    calculate_mro(new_typeinfo)
+    new_typeinfo.calculate_metaclass_type()
+
+    # 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)
+
+    classdef.info = new_typeinfo
+    current_module.names[new_class_unique_name] = SymbolTableNode(GDEF, new_typeinfo, plugin_generated=True)
+    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)

mypy_django_plugin/lib/generics.py

@@ -0,0 +1,6 @@
+def make_classes_generic(*klasses: type) -> None:
+    for klass in klasses:
+        def fake_classgetitem(cls, *args, **kwargs):
+            return cls
+
+        klass.__class_getitem__ = classmethod(fake_classgetitem)  # type: ignore