Merge pull request #427 from kszmigiel/feature/manager

Added dynamic class hook for `from_queryset` manager

Author: mkurnikov

django-stubs/contrib/admin/options.pyi

@@ -135,7 +135,7 @@ class ModelAdmin(BaseModelAdmin):
     delete_selected_confirmation_template: str = ...
     object_history_template: str = ...
     popup_response_template: str = ...
-    actions: Sequence[Callable[[ModelAdmin, HttpRequest, QuerySet], None]] = ...
+    actions: Sequence[Union[Callable[[ModelAdmin, HttpRequest, QuerySet], None], str]] = ...
     action_form: Any = ...
     actions_on_top: bool = ...
     actions_on_bottom: bool = ...

django-stubs/core/paginator.pyi

@@ -29,8 +29,8 @@ class Paginator:
         orphans: int = ...,
         allow_empty_first_page: bool = ...,
     ) -> None: ...
-    def validate_number(self, number: Optional[Union[float, str]]) -> int: ...
-    def get_page(self, number: Optional[int]) -> Page: ...
+    def validate_number(self, number: Optional[Union[int, float, str]]) -> int: ...
+    def get_page(self, number: Optional[Union[int, float, str]]) -> Page: ...
     def page(self, number: Union[int, str]) -> Page: ...
     @property
     def count(self) -> int: ...

django-stubs/db/models/enums.pyi

@@ -10,6 +10,10 @@ class ChoicesMeta(enum.EnumMeta):
 
 class Choices(enum.Enum, metaclass=ChoicesMeta):
     def __str__(self): ...
+    @property
+    def label(self) -> str: ...
+    @property
+    def value(self) -> Any: ...
 
 # fake
 class _IntegerChoicesMeta(ChoicesMeta):
@@ -18,7 +22,9 @@ class _IntegerChoicesMeta(ChoicesMeta):
     labels: List[str] = ...
     values: List[int] = ...
 
-class IntegerChoices(int, Choices, metaclass=_IntegerChoicesMeta): ...
+class IntegerChoices(int, Choices, metaclass=_IntegerChoicesMeta):
+    @property
+    def value(self) -> int: ...
 
 # fake
 class _TextChoicesMeta(ChoicesMeta):
@@ -27,4 +33,6 @@ class _TextChoicesMeta(ChoicesMeta):
     labels: List[str] = ...
     values: List[str] = ...
 
-class TextChoices(str, Choices, metaclass=_TextChoicesMeta): ...
+class TextChoices(str, Choices, metaclass=_TextChoicesMeta):
+    @property
+    def value(self) -> str: ...

django-stubs/forms/fields.pyi

@@ -14,7 +14,7 @@ class Field:
     initial: Any
     label: Optional[str]
     required: bool
-    widget: Type[Widget] = ...
+    widget: Union[Type[Widget], Widget] = ...
     hidden_widget: Any = ...
     default_validators: Any = ...
     default_error_messages: Any = ...

django-stubs/http/response.pyi

@@ -83,12 +83,13 @@ class HttpResponse(HttpResponseBase):
     context: Context
     resolver_match: ResolverMatch
     def json(self) -> Any: ...
+    def getvalue(self) -> bytes: ...
 
 class StreamingHttpResponse(HttpResponseBase):
     content: Any
     streaming_content: Iterator[Any]
     def __init__(self, streaming_content: Iterable[Any] = ..., *args: Any, **kwargs: Any) -> None: ...
-    def getvalue(self) -> Any: ...
+    def getvalue(self) -> bytes: ...
 
 class FileResponse(StreamingHttpResponse):
     client: Client

django-stubs/utils/functional.pyi

@@ -56,4 +56,8 @@ class SimpleLazyObject(LazyObject):
     def __copy__(self) -> List[int]: ...
     def __deepcopy__(self, memo: Dict[Any, Any]) -> List[int]: ...
 
-def partition(predicate: Callable, values: List[Model]) -> Tuple[List[Model], List[Model]]: ...
+_PartitionMember = TypeVar("_PartitionMember")
+
+def partition(
+    predicate: Callable, values: List[_PartitionMember]
+) -> Tuple[List[_PartitionMember], List[_PartitionMember]]: ...

mypy_django_plugin/lib/chk_helpers.py

@@ -2,8 +2,9 @@
 
 from mypy import checker
 from mypy.checker import TypeChecker
+from mypy.mro import calculate_mro
 from mypy.nodes import (
-    GDEF, MDEF, Expression, MypyFile, SymbolTableNode, TypeInfo, Var,
+    GDEF, MDEF, Block, ClassDef, Expression, MypyFile, SymbolTable, SymbolTableNode, TypeInfo, Var,
 )
 from mypy.plugin import (
     AttributeContext, CheckerPluginInterface, FunctionContext, MethodContext,
@@ -21,9 +22,17 @@ def add_new_class_for_current_module(current_module: MypyFile,
                                      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)
+
+    # 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():
@@ -32,8 +41,8 @@ def add_new_class_for_current_module(current_module: MypyFile,
             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)
-    current_module.defs.append(new_typeinfo.defn)
     return new_typeinfo
 
 

mypy_django_plugin/lib/helpers.py

@@ -6,14 +6,15 @@
 from mypy.checker import TypeChecker
 from mypy.mro import calculate_mro
 from mypy.nodes import (
-    Block, CallExpr, ClassDef, Context, Expression, MemberExpr, MypyFile, NameExpr, PlaceholderNode, StrExpr,
-    SymbolTable, SymbolTableNode, TypeInfo, Var,
+    GDEF, Argument, Block, CallExpr, ClassDef, Context, Expression, FuncDef, MemberExpr, MypyFile, NameExpr,
+    PlaceholderNode, StrExpr, SymbolTable, SymbolTableNode, TypeInfo, Var,
 )
 from mypy.plugin import (
     AttributeContext, ClassDefContext, DynamicClassDefContext, FunctionContext, MethodContext,
 )
+from mypy.plugins.common import add_method_to_class
 from mypy.semanal import SemanticAnalyzer
-from mypy.types import AnyType, Instance, NoneTyp, ProperType
+from mypy.types import AnyType, CallableType, Instance, NoneTyp, ProperType
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, UnionType
 
@@ -94,6 +95,80 @@ def lookup_typeinfo_or_defer(self, fullname: str, *,
 
         return sym.node
 
+    def copy_method_to_another_class(
+            self,
+            ctx: ClassDefContext,
+            self_type: Instance,
+            new_method_name: str,
+            method_node: FuncDef) -> None:
+        if method_node.type is None:
+            if not self.defer_till_next_iteration(reason='method_node.type is None'):
+                raise new_helpers.TypeInfoNotFound(method_node.fullname)
+
+            arguments, return_type = build_unannotated_method_args(method_node)
+            add_method_to_class(
+                ctx.api,
+                ctx.cls,
+                new_method_name,
+                args=arguments,
+                return_type=return_type,
+                self_type=self_type)
+            return
+
+        method_type = cast(CallableType, method_node.type)
+        if not isinstance(method_type, CallableType) and not self.defer_till_next_iteration(
+                reason='method_node.type is not CallableType'):
+            raise new_helpers.TypeInfoNotFound(method_node.fullname)
+
+        arguments = []
+        bound_return_type = self.semanal_api.anal_type(
+            method_type.ret_type,
+            allow_placeholder=True)
+
+        if bound_return_type is None and self.defer_till_next_iteration():
+            raise new_helpers.TypeInfoNotFound(method_node.fullname + ' return type')
+
+        assert bound_return_type is not None
+
+        if isinstance(bound_return_type, PlaceholderNode):
+            raise new_helpers.TypeInfoNotFound('return type ' + method_node.fullname)
+
+        for arg_name, arg_type, original_argument in zip(
+                method_type.arg_names[1:],
+                method_type.arg_types[1:],
+                method_node.arguments[1:]):
+            bound_arg_type = self.semanal_api.anal_type(arg_type, allow_placeholder=True)
+            if bound_arg_type is None and not self.defer_till_next_iteration(reason='bound_arg_type is None'):
+                error_msg = 'of {} argument of {}'.format(arg_name, method_node.fullname)
+                raise new_helpers.TypeInfoNotFound(error_msg)
+
+            assert bound_arg_type is not None
+
+            if isinstance(bound_arg_type, PlaceholderNode) and self.defer_till_next_iteration(
+                    reason='bound_arg_type is None'):
+                raise new_helpers.TypeInfoNotFound('of ' + arg_name + ' argument of ' + method_node.fullname)
+
+            var = Var(
+                name=original_argument.variable.name,
+                type=arg_type)
+            var.line = original_argument.variable.line
+            var.column = original_argument.variable.column
+            argument = Argument(
+                variable=var,
+                type_annotation=bound_arg_type,
+                initializer=original_argument.initializer,
+                kind=original_argument.kind)
+            argument.set_line(original_argument)
+            arguments.append(argument)
+
+        add_method_to_class(
+            ctx.api,
+            ctx.cls,
+            new_method_name,
+            args=arguments,
+            return_type=bound_return_type,
+            self_type=self_type)
+
     def new_typeinfo(self, name: str, bases: List[Instance], module_fullname: Optional[str] = None) -> TypeInfo:
         class_def = ClassDef(name, Block([]))
         class_def.fullname = self.semanal_api.qualified_name(name)
@@ -118,11 +193,46 @@ def __call__(self, ctx: DynamicClassDefContext) -> None:
         self.semanal_api = cast(SemanticAnalyzer, ctx.api)
         self.create_new_dynamic_class()
 
+    def generate_manager_info_and_module(self, base_manager_info: TypeInfo) -> Tuple[TypeInfo, MypyFile]:
+        new_manager_info = self.semanal_api.basic_new_typeinfo(
+            self.class_name,
+            basetype_or_fallback=Instance(
+                base_manager_info,
+                [AnyType(TypeOfAny.unannotated)])
+        )
+        new_manager_info.line = self.call_expr.line
+        new_manager_info.defn.line = self.call_expr.line
+        new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type()
+
+        current_module = self.semanal_api.cur_mod_node
+        current_module.names[self.class_name] = SymbolTableNode(
+            GDEF,
+            new_manager_info,
+            plugin_generated=True)
+        return new_manager_info, current_module
+
     @abstractmethod
     def create_new_dynamic_class(self) -> None:
         raise NotImplementedError
 
 
+class DynamicClassFromMethodCallback(DynamicClassPluginCallback):
+    callee: MemberExpr
+
+    def __call__(self, ctx: DynamicClassDefContext) -> None:
+        self.class_name = ctx.name
+        self.call_expr = ctx.call
+
+        assert ctx.call.callee is not None
+        if not isinstance(ctx.call.callee, MemberExpr):
+            # throw error?
+            return
+        self.callee = ctx.call.callee
+
+        self.semanal_api = cast(SemanticAnalyzer, ctx.api)
+        self.create_new_dynamic_class()
+
+
 class ClassDefPluginCallback(SemanalPluginCallback):
     reason: Expression
     class_defn: ClassDef
@@ -396,3 +506,12 @@ def get_nested_meta_node_for_current_class(info: TypeInfo) -> Optional[TypeInfo]
     if metaclass_sym is not None and isinstance(metaclass_sym.node, TypeInfo):
         return metaclass_sym.node
     return None
+
+
+def build_unannotated_method_args(method_node: FuncDef) -> Tuple[List[Argument], MypyType]:
+    prepared_arguments = []
+    for argument in method_node.arguments[1:]:
+        argument.type_annotation = AnyType(TypeOfAny.unannotated)
+        prepared_arguments.append(argument)
+    return_type = AnyType(TypeOfAny.unannotated)
+    return prepared_arguments, return_type

mypy_django_plugin/main.py

@@ -6,7 +6,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
 
@@ -19,6 +19,9 @@
 from mypy_django_plugin.transformers.init_create import (
     ModelCreateCallback, ModelInitCallback,
 )
+from mypy_django_plugin.transformers.managers import (
+    ManagerFromQuerySetCallback,
+)
 from mypy_django_plugin.transformers.meta import MetaGetFieldCallback
 from mypy_django_plugin.transformers.models import ModelCallback
 from mypy_django_plugin.transformers.orm_lookups import (
@@ -249,6 +252,15 @@ def get_attribute_hook(self, fullname: str
 
         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 ManagerFromQuerySetCallback(self)
+        return None
+
 
 def plugin(version):
     return NewSemanalDjangoPlugin

mypy_django_plugin/transformers/managers.py

@@ -0,0 +1,76 @@
+from typing import List, Tuple
+
+from mypy.nodes import Argument, FuncDef, NameExpr, StrExpr, TypeInfo
+from mypy.plugin import ClassDefContext
+from mypy.types import AnyType, Instance
+from mypy.types import Type as MypyType
+from mypy.types import TypeOfAny
+
+from mypy_django_plugin.lib import fullnames, helpers
+
+
+def build_unannotated_method_args(method_node: FuncDef) -> Tuple[List[Argument], MypyType]:
+    prepared_arguments = []
+    for argument in method_node.arguments[1:]:
+        argument.type_annotation = AnyType(TypeOfAny.unannotated)
+        prepared_arguments.append(argument)
+    return_type = AnyType(TypeOfAny.unannotated)
+    return prepared_arguments, return_type
+
+
+class ManagerFromQuerySetCallback(helpers.DynamicClassFromMethodCallback):
+    def create_new_dynamic_class(self) -> None:
+
+        base_manager_info = self.callee.expr.node  # type: ignore
+
+        if base_manager_info is None and not self.defer_till_next_iteration(reason='base_manager_info is None'):
+            # what exception should be thrown here?
+            return
+
+        assert isinstance(base_manager_info, TypeInfo)
+
+        new_manager_info, current_module = self.generate_manager_info_and_module(base_manager_info)
+
+        passed_queryset = self.call_expr.args[0]
+        assert isinstance(passed_queryset, NameExpr)
+
+        derived_queryset_fullname = passed_queryset.fullname
+        assert derived_queryset_fullname is not None
+
+        sym = self.semanal_api.lookup_fully_qualified_or_none(derived_queryset_fullname)
+        assert sym is not None
+        if sym.node is None and not self.defer_till_next_iteration(reason='sym.node is None'):
+            # 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(self.call_expr.args) > 1:
+            expr = self.call_expr.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=self.call_expr, api=self.semanal_api)
+        self_type = Instance(new_manager_info, [])
+        # we need to copy all methods in MRO before django.db.models.query.QuerySet
+        for class_mro_info in derived_queryset_info.mro:
+            if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME:
+                break
+            for name, sym in class_mro_info.names.items():
+                if isinstance(sym.node, FuncDef):
+                    self.copy_method_to_another_class(
+                        class_def_context,
+                        self_type,
+                        new_method_name=name,
+                        method_node=sym.node)