This commit is contained in:
Maxim Kurnikov
2020-02-02 03:12:32 +03:00
parent a01d58462e
commit 0b1507c81e
17 changed files with 847 additions and 363 deletions

View File

@@ -55,10 +55,10 @@ def initialize_django(settings_module: str) -> Tuple['Apps', 'LazySettings']:
def noop_class_getitem(cls, key):
return cls
from django.db import models
# 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
# 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
@@ -226,7 +226,7 @@ class DjangoContext:
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

View File

@@ -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

View File

@@ -1,8 +1,8 @@
from abc import abstractmethod
from typing import (
TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union,
)
cast)
from django.db.models.fields import Field
from django.db.models.fields.related import RelatedField
from django.db.models.fields.reverse_related import ForeignObjectRel
from mypy.checker import TypeChecker
@@ -10,20 +10,210 @@ from mypy.mro import calculate_mro
from mypy.nodes import (
Block, ClassDef, Expression, MemberExpr, MypyFile, NameExpr, StrExpr, SymbolTable, SymbolTableNode,
TypeInfo, Var,
)
CallExpr, Context, PlaceholderNode, FuncDef, FakeInfo)
from mypy.plugin import DynamicClassDefContext, ClassDefContext
from mypy.plugins.common import add_method
from mypy.semanal import SemanticAnalyzer
from mypy.types import AnyType, Instance, NoneTyp
from mypy.types import AnyType, Instance, NoneTyp, TypeType
from mypy.types import Type as MypyType
from mypy.types import TypeOfAny, UnionType
from mypy.typetraverser import TypeTraverserVisitor
from django.db.models.fields import Field
from mypy_django_plugin.lib import fullnames
from mypy_django_plugin.lib.sem_helpers import prepare_unannotated_method_signature, analyze_callable_signature
from mypy_django_plugin.transformers2 import new_helpers
if TYPE_CHECKING:
from mypy_django_plugin.django.context import DjangoContext
from mypy_django_plugin.main import NewSemanalDjangoPlugin
AnyPluginAPI = Union[TypeChecker, SemanticAnalyzer]
class DjangoPluginCallback:
django_context: 'DjangoContext'
def __init__(self, plugin: 'NewSemanalDjangoPlugin') -> None:
self.plugin = plugin
self.django_context = plugin.django_context
# def lookup_fully_qualified(self, fullname: str) -> Optional[SymbolTableNode]:
# return self.plugin.lookup_fully_qualified(fullname)
class SemanalPluginCallback(DjangoPluginCallback):
semanal_api: SemanticAnalyzer
def build_defer_error_message(self, message: str) -> str:
return f'{self.__class__.__name__}: {message}'
def defer_till_next_iteration(self, deferral_context: Optional[Context] = None,
*,
reason: Optional[str] = None) -> bool:
""" Returns False if cannot be deferred. """
if self.semanal_api.final_iteration:
return False
self.semanal_api.defer(deferral_context)
print(f'LOG: defer: {self.build_defer_error_message(reason)}')
return True
def lookup_typeinfo_or_defer(self, fullname: str, *,
deferral_context: Optional[Context] = None,
reason_for_defer: Optional[str] = None) -> Optional[TypeInfo]:
sym = self.plugin.lookup_fully_qualified(fullname)
if sym is None or sym.node is None or isinstance(sym.node, PlaceholderNode):
deferral_context = deferral_context or self.semanal_api.cur_mod_node
reason = reason_for_defer or f'{fullname!r} is not available for lookup'
if not self.defer_till_next_iteration(deferral_context, reason=reason):
raise new_helpers.TypeInfoNotFound(fullname)
return None
if not isinstance(sym.node, TypeInfo):
raise ValueError(f'{fullname!r} does not correspond to TypeInfo')
return sym.node
def new_typeinfo(self, name: str, bases: List[Instance]) -> TypeInfo:
class_def = ClassDef(name, Block([]))
class_def.fullname = self.semanal_api.qualified_name(name)
info = TypeInfo(SymbolTable(), class_def, self.semanal_api.cur_mod_id)
info.bases = bases
calculate_mro(info)
info.metaclass_type = info.calculate_metaclass_type()
class_def.info = info
return info
# def add_symbol_table_node_or_defer(self, name: str, sym: SymbolTableNode) -> bool:
# return self.semanal_api.add_symbol_table_node(name, sym,
# context=self.semanal_api.cur_mod_node)
def add_method_from_signature(self,
signature_node: FuncDef,
new_method_name: str,
new_self_type: Instance,
class_defn: ClassDef) -> bool:
if signature_node.type is None:
if self.defer_till_next_iteration(reason=signature_node.fullname):
return False
arguments, return_type = prepare_unannotated_method_signature(signature_node)
ctx = ClassDefContext(class_defn, signature_node, self.semanal_api)
add_method(ctx,
new_method_name,
self_type=new_self_type,
args=arguments,
return_type=return_type)
return True
# add imported objects from method signature to the current module, if not present
source_symbols = self.semanal_api.modules[signature_node.info.module_name].names
currently_imported_symbols = self.semanal_api.cur_mod_node.names
def import_symbol_from_source(name: str) -> None:
if name in source_symbols['__builtins__'].node.names:
return
sym = source_symbols[name].copy()
self.semanal_api.add_imported_symbol(name, sym, context=self.semanal_api.cur_mod_node)
class UnimportedTypesVisitor(TypeTraverserVisitor):
def visit_union_type(self, t: UnionType) -> None:
super().visit_union_type(t)
union_sym = currently_imported_symbols.get('Union')
if union_sym is None:
# TODO: check if it's exactly typing.Union
import_symbol_from_source('Union')
def visit_type_type(self, t: TypeType) -> None:
super().visit_type_type(t)
type_sym = currently_imported_symbols.get('Union')
if type_sym is None:
# TODO: check if it's exactly typing.Type
import_symbol_from_source('Type')
def visit_instance(self, t: Instance) -> None:
super().visit_instance(t)
if isinstance(t.type, FakeInfo):
return
type_name = t.type.name
sym = currently_imported_symbols.get(type_name)
if sym is None:
# TODO: check if it's exactly typing.Type
import_symbol_from_source(type_name)
signature_node.type.accept(UnimportedTypesVisitor())
# # copy global SymbolTableNode objects from original class to the current node, if not present
# original_module = semanal_api.modules[method_node.info.module_name]
# for name, sym in original_module.names.items():
# if (not sym.plugin_generated
# and name not in semanal_api.cur_mod_node.names):
# semanal_api.add_imported_symbol(name, sym, context=semanal_api.cur_mod_node)
arguments, analyzed_return_type, unbound = analyze_callable_signature(self.semanal_api, signature_node)
if unbound:
raise new_helpers.IncompleteDefnError(f'Signature of method {signature_node.fullname!r} is not ready')
assert len(arguments) + 1 == len(signature_node.arguments)
assert analyzed_return_type is not None
ctx = ClassDefContext(class_defn, signature_node, self.semanal_api)
add_method(ctx,
new_method_name,
self_type=new_self_type,
args=arguments,
return_type=analyzed_return_type)
return True
class DynamicClassPluginCallback(SemanalPluginCallback):
class_name: str
call_expr: CallExpr
def __call__(self, ctx: DynamicClassDefContext) -> None:
self.class_name = ctx.name
self.call_expr = ctx.call
self.semanal_api = cast(SemanticAnalyzer, ctx.api)
self.create_new_dynamic_class()
def get_callee(self) -> MemberExpr:
callee = self.call_expr.callee
assert isinstance(callee, MemberExpr)
return callee
def lookup_same_module_or_defer(self, name: str, *,
deferral_context: Optional[Context] = None) -> Optional[SymbolTableNode]:
sym = self.semanal_api.lookup_qualified(name, self.call_expr)
if sym is None or sym.node is None or isinstance(sym.node, PlaceholderNode):
deferral_context = deferral_context or self.call_expr
if not self.defer_till_next_iteration(deferral_context,
reason=f'{self.semanal_api.cur_mod_id}.{name} does not exist'):
raise new_helpers.NameNotFound(name)
return None
return sym
@abstractmethod
def create_new_dynamic_class(self) -> None:
raise NotImplementedError
class ClassDefPluginCallback(SemanalPluginCallback):
reason: Expression
class_defn: ClassDef
def __call__(self, ctx: ClassDefContext) -> None:
self.reason = ctx.reason
self.class_defn = ctx.cls
self.semanal_api = cast(SemanticAnalyzer, ctx.api)
self.modify_class_defn()
@abstractmethod
def modify_class_defn(self) -> None:
raise NotImplementedError
def get_django_metadata(model_info: TypeInfo) -> Dict[str, Any]:
return model_info.metadata.setdefault('django', {})
@@ -31,7 +221,6 @@ def get_django_metadata(model_info: TypeInfo) -> Dict[str, Any]:
def split_symbol_name(fullname: str, all_modules: Dict[str, MypyFile]) -> Optional[Tuple[str, str]]:
if '.' not in fullname:
return None
module_name = None
parts = fullname.split('.')
for i in range(len(parts), 0, -1):
@@ -39,12 +228,11 @@ def split_symbol_name(fullname: str, all_modules: Dict[str, MypyFile]) -> Option
if possible_module_name in all_modules:
module_name = possible_module_name
break
if module_name is None:
return None
cls_name = fullname.replace(module_name, '').lstrip('.')
return module_name, cls_name
symbol_name = fullname.replace(module_name, '').lstrip('.')
return module_name, symbol_name
def lookup_fully_qualified_typeinfo(api: AnyPluginAPI, fullname: str) -> Optional[TypeInfo]:

View File

@@ -9,12 +9,12 @@ from mypy.types import Type as MypyType
from mypy.types import TypeOfAny, get_proper_type
class IncompleteDefnException(Exception):
class IncompleteDefnError(Exception):
def __init__(self, error_message: str = '') -> None:
super().__init__(error_message)
class BoundNameNotFound(IncompleteDefnException):
class BoundNameNotFound(IncompleteDefnError):
def __init__(self, fullname: str) -> None:
super().__init__(f'No {fullname!r} found')
@@ -85,7 +85,7 @@ def copy_method_or_incomplete_defn_exception(ctx: ClassDefContext,
if method_node.type is None:
if not semanal_api.final_iteration:
raise IncompleteDefnException(f'Unannotated method {method_node.fullname!r}')
raise IncompleteDefnError(f'Unannotated method {method_node.fullname!r}')
arguments, return_type = prepare_unannotated_method_signature(method_node)
add_method(ctx,
@@ -107,7 +107,7 @@ def copy_method_or_incomplete_defn_exception(ctx: ClassDefContext,
arguments, analyzed_return_type, unbound = analyze_callable_signature(semanal_api, method_node)
assert len(arguments) + 1 == len(method_node.arguments)
if unbound:
raise IncompleteDefnException(f'Signature of method {method_node.fullname!r} is not ready')
raise IncompleteDefnError(f'Signature of method {method_node.fullname!r} is not ready')
assert analyzed_return_type is not None

View File

@@ -18,10 +18,10 @@ from mypy_django_plugin.transformers import (
fields, forms, init_create, meta, querysets, request, settings,
)
from mypy_django_plugin.transformers.managers import (
create_manager_class_from_as_manager_method, create_new_manager_class_from_from_queryset_method,
instantiate_anonymous_queryset_from_as_manager,
)
create_manager_class_from_as_manager_method, instantiate_anonymous_queryset_from_as_manager)
from mypy_django_plugin.transformers.models import process_model_class
from mypy_django_plugin.transformers2.dynamic_managers import CreateNewManagerClassFrom_FromQuerySet
from mypy_django_plugin.transformers2.models import ModelCallback
def transform_model_class(ctx: ClassDefContext,
@@ -176,9 +176,9 @@ class NewSemanalDjangoPlugin(Plugin):
if fullname == 'django.contrib.auth.get_user_model':
return partial(settings.get_user_model_hook, django_context=self.django_context)
manager_bases = self._get_current_manager_bases()
if fullname in manager_bases:
return querysets.determine_proper_manager_type
# manager_bases = self._get_current_manager_bases()
# if fullname in manager_bases:
# return querysets.determine_proper_manager_type
info = self._get_typeinfo_or_none(fullname)
if info:
@@ -234,7 +234,7 @@ class NewSemanalDjangoPlugin(Plugin):
) -> Optional[Callable[[ClassDefContext], None]]:
if (fullname in self.django_context.all_registered_model_class_fullnames
or fullname in self._get_current_model_bases()):
return partial(transform_model_class, django_context=self.django_context)
return ModelCallback(self)
if fullname in self._get_current_manager_bases():
return add_new_manager_base
@@ -261,7 +261,8 @@ class NewSemanalDjangoPlugin(Plugin):
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 CreateNewManagerClassFrom_FromQuerySet(self)
if fullname.endswith('as_manager'):
class_name, _, _ = fullname.rpartition('.')
info = self._get_typeinfo_or_none(class_name)

View File

View File

@@ -110,6 +110,11 @@ def get_field_descriptor_types(field_info: TypeInfo, is_nullable: bool) -> Tuple
return set_type, get_type
def get_field_type(field_info: TypeInfo, is_nullable: bool) -> Instance:
set_type, get_type = get_field_descriptor_types(field_info, is_nullable)
return Instance(field_info, [set_type, get_type])
def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance:
default_return_type = cast(Instance, ctx.default_return_type)

View File

@@ -1,12 +1,13 @@
from typing import Any, Dict, Iterator, Optional, Tuple
from mypy.checker import gen_unique_name
from mypy.nodes import (
GDEF, CallExpr, Context, Decorator, FuncDef, MemberExpr, NameExpr, OverloadedFuncDef, PlaceholderNode, RefExpr,
StrExpr, SymbolTable, SymbolTableNode, TypeInfo,
)
MypyFile)
from mypy.plugin import ClassDefContext, DynamicClassDefContext, MethodContext
from mypy.semanal import SemanticAnalyzer, is_same_symbol, is_valid_replacement
from mypy.types import AnyType, CallableType, Instance
from mypy.types import AnyType, CallableType, Instance, TypeVarType, TypeVarDef
from mypy.types import Type as MypyType
from mypy.types import TypeOfAny
@@ -26,28 +27,43 @@ def generate_from_queryset_name(base_manager_info: TypeInfo, queryset_info: Type
return base_manager_info.name + 'From' + queryset_info.name
def resolve_callee_info_or_exception(ctx: DynamicClassDefContext) -> TypeInfo:
callee = ctx.call.callee
#
# def cb_resolve_callee_info_or_exception(cb: ) -> TypeInfo:
# callee = ctx.call.callee
# assert isinstance(callee, MemberExpr)
# assert isinstance(callee.expr, RefExpr)
#
# callee_info = callee.expr.node
# if (callee_info is None
# or isinstance(callee_info, PlaceholderNode)):
# raise sem_helpers.IncompleteDefnError(f'Definition of base manager {callee.fullname!r} '
# f'is incomplete.')
#
# assert isinstance(callee_info, TypeInfo)
# return callee_info
def resolve_callee_info_or_exception(callback: helpers.DynamicClassPluginCallback) -> TypeInfo:
callee = callback.call_expr.callee
assert isinstance(callee, MemberExpr)
assert isinstance(callee.expr, RefExpr)
callee_info = callee.expr.node
if (callee_info is None
or isinstance(callee_info, PlaceholderNode)):
raise sem_helpers.IncompleteDefnException(f'Definition of base manager {callee.fullname!r} '
f'is incomplete.')
raise sem_helpers.IncompleteDefnError(f'Definition of base manager {callee.fullname!r} '
f'is incomplete.')
assert isinstance(callee_info, TypeInfo)
return callee_info
def resolve_passed_queryset_info_or_exception(ctx: DynamicClassDefContext) -> TypeInfo:
api = sem_helpers.get_semanal_api(ctx)
passed_queryset_name_expr = ctx.call.args[0]
def resolve_passed_queryset_info_or_exception(callback: helpers.DynamicClassPluginCallback) -> TypeInfo:
passed_queryset_name_expr = callback.call_expr.args[0]
assert isinstance(passed_queryset_name_expr, NameExpr)
sym = api.lookup_qualified(passed_queryset_name_expr.name, ctx=ctx.call)
# lookup in the same module
sym = callback.semanal_api.lookup_qualified(passed_queryset_name_expr.name, ctx=callback.call_expr)
if (sym is None
or sym.node is None
or isinstance(sym.node, PlaceholderNode)):
@@ -58,9 +74,8 @@ def resolve_passed_queryset_info_or_exception(ctx: DynamicClassDefContext) -> Ty
return sym.node
def resolve_django_manager_info_or_exception(ctx: DynamicClassDefContext) -> TypeInfo:
api = sem_helpers.get_semanal_api(ctx)
info = helpers.lookup_fully_qualified_typeinfo(api, fullnames.MANAGER_CLASS_FULLNAME)
def resolve_django_manager_info_or_exception(callback: helpers.DynamicClassPluginCallback) -> TypeInfo:
info = callback.lookup_typeinfo_or_defer(fullnames.MANAGER_CLASS_FULLNAME)
if info is None:
raise sem_helpers.BoundNameNotFound(fullnames.MANAGER_CLASS_FULLNAME)
@@ -113,7 +128,7 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
callee_manager_info = resolve_callee_info_or_exception(ctx)
queryset_info = resolve_passed_queryset_info_or_exception(ctx)
django_manager_info = resolve_django_manager_info_or_exception(ctx)
except sem_helpers.IncompleteDefnException:
except sem_helpers.IncompleteDefnError:
if not semanal_api.final_iteration:
semanal_api.defer()
return
@@ -137,7 +152,7 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
self_type,
new_method_name=name,
method_node=method_node)
except sem_helpers.IncompleteDefnException:
except sem_helpers.IncompleteDefnError:
if not semanal_api.final_iteration:
semanal_api.defer()
return
@@ -216,12 +231,17 @@ def add_symbol_table_node(api: SemanticAnalyzer,
return False
class CreateNewManagerClassFrom_AsManager(helpers.DynamicClassPluginCallback):
def create_new_dynamic_class(self) -> None:
pass
def create_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) -> None:
semanal_api = sem_helpers.get_semanal_api(ctx)
try:
queryset_info = resolve_callee_info_or_exception(ctx)
django_manager_info = resolve_django_manager_info_or_exception(ctx)
except sem_helpers.IncompleteDefnException:
except sem_helpers.IncompleteDefnError:
if not semanal_api.final_iteration:
semanal_api.defer()
return
@@ -258,7 +278,7 @@ def create_manager_class_from_as_manager_method(ctx: DynamicClassDefContext) ->
self_type,
new_method_name=name,
method_node=method_node)
except sem_helpers.IncompleteDefnException:
except sem_helpers.IncompleteDefnError:
if not semanal_api.final_iteration:
semanal_api.defer()
return

View File

@@ -7,7 +7,7 @@ from django.db.models.fields.reverse_related import (
ManyToManyRel, ManyToOneRel, OneToOneRel,
)
from mypy.nodes import (
ARG_STAR2, GDEF, MDEF, Argument, Context, FuncDef, SymbolTableNode, TypeInfo, Var,
ARG_STAR2, GDEF, MDEF, Argument, Context, SymbolTableNode, TypeInfo, Var,
)
from mypy.plugin import ClassDefContext
from mypy.plugins import common
@@ -34,15 +34,15 @@ class ModelClassInitializer:
def lookup_typeinfo(self, fullname: str) -> Optional[TypeInfo]:
return helpers.lookup_fully_qualified_typeinfo(self.api, fullname)
def lookup_typeinfo_or_incomplete_defn_error(self, fullname: str) -> TypeInfo:
def lookup_typeinfo_or_exception(self, fullname: str) -> TypeInfo:
info = self.lookup_typeinfo(fullname)
if info is None:
raise sem_helpers.IncompleteDefnException(f'No {fullname!r} found')
raise sem_helpers.IncompleteDefnError(f'No {fullname!r} found')
return info
def lookup_class_typeinfo_or_incomplete_defn_error(self, klass: type) -> TypeInfo:
def lookup_class_typeinfo_or_exception(self, klass: type) -> TypeInfo:
fullname = helpers.get_class_fullname(klass)
field_info = self.lookup_typeinfo_or_incomplete_defn_error(fullname)
field_info = self.lookup_typeinfo_or_exception(fullname)
return field_info
def model_class_has_attribute_defined(self, name: str, traverse_mro: bool = True) -> bool:
@@ -141,7 +141,7 @@ class AddDefaultPrimaryKey(ModelClassInitializer):
return
auto_field_class_fullname = helpers.get_class_fullname(auto_field.__class__)
auto_field_info = self.lookup_typeinfo_or_incomplete_defn_error(auto_field_class_fullname)
auto_field_info = self.lookup_typeinfo_or_exception(auto_field_class_fullname)
set_type, get_type = fields.get_field_descriptor_types(auto_field_info, is_nullable=False)
self.add_new_node_to_model_class(primary_key_attrname, Instance(auto_field_info,
@@ -168,8 +168,6 @@ class AddRelatedModelsId(ModelClassInitializer):
related_id_attr_name = field.attname
if self.model_class_has_attribute_defined(related_id_attr_name):
continue
# if self.get_model_class_attr(related_id_attr_name) is not None:
# continue
related_model_cls = self.django_context.get_field_related_model_cls(field)
if related_model_cls is None:
@@ -189,8 +187,8 @@ class AddRelatedModelsId(ModelClassInitializer):
rel_primary_key_field = self.django_context.get_primary_key_field(related_model_cls)
try:
field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_primary_key_field.__class__)
except sem_helpers.IncompleteDefnException as exc:
field_info = self.lookup_class_typeinfo_or_exception(rel_primary_key_field.__class__)
except sem_helpers.IncompleteDefnError as exc:
if not self.api.final_iteration:
raise exc
else:
@@ -203,53 +201,6 @@ class AddRelatedModelsId(ModelClassInitializer):
class AddManagers(ModelClassInitializer):
def has_any_parametrized_manager_as_base(self, info: TypeInfo) -> bool:
for base in helpers.iter_bases(info):
if self.is_any_parametrized_manager(base):
return True
return False
def is_any_parametrized_manager(self, typ: Instance) -> bool:
return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType)
def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance:
bases = []
for original_base in base_manager_info.bases:
if self.is_any_parametrized_manager(original_base):
original_base = helpers.reparametrize_instance(original_base,
[Instance(self.model_classdef.info, [])])
bases.append(original_base)
new_manager_info = self.add_new_class_for_current_module(name, bases, force_replace_existing=True)
# copy fields to a new manager
new_cls_def_context = ClassDefContext(cls=new_manager_info.defn,
reason=self.ctx.reason,
api=self.api)
custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])])
for name, sym in base_manager_info.names.items():
if name in new_manager_info.names:
raise ValueError(f'Name {name!r} already exists on newly-created {new_manager_info.fullname!r} class.')
# replace self type with new class, if copying method
if isinstance(sym.node, FuncDef):
sem_helpers.copy_method_or_incomplete_defn_exception(new_cls_def_context,
self_type=custom_manager_type,
new_method_name=name,
method_node=sym.node)
continue
new_sym = sym.copy()
if isinstance(new_sym.node, Var):
new_var = Var(name, type=sym.type)
new_var.info = new_manager_info
new_var._fullname = new_manager_info.fullname + '.' + name
new_sym.node = new_var
new_manager_info.names[name] = new_sym
return custom_manager_type
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
for manager_name, manager in model_cls._meta.managers_map.items():
if self.model_class_has_attribute_defined(manager_name, traverse_mro=False):
@@ -258,8 +209,7 @@ class AddManagers(ModelClassInitializer):
if (sym.type is not None
and isinstance(sym.type, Instance)
and sym.type.type.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
and not self.has_any_parametrized_manager_as_base(sym.type.type)):
and sym.type.type.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)):
# already defined and parametrized properly
continue
@@ -270,24 +220,12 @@ class AddManagers(ModelClassInitializer):
continue
manager_fullname = self.resolve_manager_fullname(helpers.get_class_fullname(manager.__class__))
manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)
manager_class_name = manager_fullname.rsplit('.', maxsplit=1)[1]
manager_info = self.lookup_typeinfo_or_exception(manager_fullname)
if manager_name not in self.model_classdef.info.names:
# manager not yet defined, just add models.Manager[ModelName]
manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
self.add_new_node_to_model_class(manager_name, manager_type)
else:
# creates new MODELNAME_MANAGERCLASSNAME class that represents manager parametrized with current model
if not self.has_any_parametrized_manager_as_base(manager_info):
continue
custom_model_manager_name = manager.model.__name__ + '_' + manager_class_name
custom_manager_type = self.create_new_model_parametrized_manager(custom_model_manager_name,
base_manager_info=manager_info)
self.add_new_node_to_model_class(manager_name, custom_manager_type,
force_replace_existing=True)
class AddDefaultManagerAttribute(ModelClassInitializer):
@@ -304,7 +242,7 @@ class AddDefaultManagerAttribute(ModelClassInitializer):
default_manager_fullname = helpers.get_class_fullname(model_cls._meta.default_manager.__class__)
resolved_default_manager_fullname = self.resolve_manager_fullname(default_manager_fullname)
default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(resolved_default_manager_fullname)
default_manager_info = self.lookup_typeinfo_or_exception(resolved_default_manager_fullname)
default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])])
self.add_new_node_to_model_class('_default_manager', default_manager)
@@ -323,11 +261,12 @@ class AddRelatedManagers(ModelClassInitializer):
related_model_cls = self.django_context.get_field_related_model_cls(relation)
if related_model_cls is None:
# could not find a referenced model (maybe invalid to= value)
continue
try:
related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(related_model_cls)
except sem_helpers.IncompleteDefnException as exc:
related_model_info = self.lookup_class_typeinfo_or_exception(related_model_cls)
except sem_helpers.IncompleteDefnError as exc:
if not self.api.final_iteration:
raise exc
else:
@@ -339,11 +278,11 @@ class AddRelatedManagers(ModelClassInitializer):
if isinstance(relation, (ManyToOneRel, ManyToManyRel)):
try:
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
related_manager_info = self.lookup_typeinfo_or_exception(
fullnames.RELATED_MANAGER_CLASS) # noqa: E501
if 'objects' not in related_model_info.names:
raise sem_helpers.IncompleteDefnException()
except sem_helpers.IncompleteDefnException as exc:
raise sem_helpers.IncompleteDefnError()
except sem_helpers.IncompleteDefnError as exc:
if not self.api.final_iteration:
raise exc
else:
@@ -372,7 +311,7 @@ class AddExtraFieldMethods(ModelClassInitializer):
# get_FOO_display for choices
for field in self.django_context.get_model_fields(model_cls):
if field.choices:
info = self.lookup_typeinfo_or_incomplete_defn_error('builtins.str')
info = self.lookup_typeinfo_or_exception('builtins.str')
return_type = Instance(info, [])
common.add_method(self.ctx,
name='get_{}_display'.format(field.attname),
@@ -402,7 +341,7 @@ class AddExtraFieldMethods(ModelClassInitializer):
class AddMetaOptionsAttribute(ModelClassInitializer):
def run_with_model_cls(self, model_cls: Type[Model]) -> None:
if '_meta' not in self.model_classdef.info.names:
options_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.OPTIONS_CLASS_FULLNAME)
options_info = self.lookup_typeinfo_or_exception(fullnames.OPTIONS_CLASS_FULLNAME)
self.add_new_node_to_model_class('_meta',
Instance(options_info, [
Instance(self.model_classdef.info, [])
@@ -424,7 +363,7 @@ def process_model_class(ctx: ClassDefContext,
for initializer_cls in initializers:
try:
initializer_cls(ctx, django_context).run()
except sem_helpers.IncompleteDefnException as exc:
except sem_helpers.IncompleteDefnError as exc:
if not ctx.api.final_iteration:
ctx.api.defer()
continue

View File

@@ -0,0 +1,79 @@
from typing import Optional
from mypy.checker import gen_unique_name
from mypy.nodes import NameExpr, TypeInfo, SymbolTableNode, StrExpr
from mypy.types import Type as MypyType, TypeVarType, TypeVarDef, Instance
from mypy_django_plugin.lib import helpers, fullnames, chk_helpers, sem_helpers
from mypy_django_plugin.transformers.managers import iter_all_custom_queryset_methods
class CreateNewManagerClassFrom_FromQuerySet(helpers.DynamicClassPluginCallback):
def create_typevar_in_current_module(self, name: str,
upper_bound: Optional[MypyType] = None) -> TypeVarDef:
tvar_name = gen_unique_name(name, self.semanal_api.globals)
tvar_def = TypeVarDef(tvar_name,
fullname=self.semanal_api.cur_mod_id + '.' + tvar_name,
id=-1,
values=[],
upper_bound=upper_bound)
return tvar_def
def create_new_dynamic_class(self) -> None:
# extract Manager class which will act as base
callee = self.get_callee()
fullname = callee.fullname or callee.expr.fullname
callee_manager_info = self.lookup_typeinfo_or_defer(fullname)
if callee_manager_info is None:
return None
# extract queryset from which we're going to copy methods
passed_queryset_name_expr = self.call_expr.args[0]
assert isinstance(passed_queryset_name_expr, NameExpr)
queryset_class_name = passed_queryset_name_expr.name
sym = self.lookup_same_module_or_defer(queryset_class_name)
if sym is None:
return None
assert isinstance(sym.node, TypeInfo)
passed_queryset_info = sym.node
# for TypeVar bound
base_model_info = self.lookup_typeinfo_or_defer(fullnames.MODEL_CLASS_FULLNAME)
if base_model_info is None:
return
model_tvar_defn = self.create_typevar_in_current_module('_M', upper_bound=Instance(base_model_info, []))
model_tvar_type = TypeVarType(model_tvar_defn)
# make Manager[_T]
parent_manager_type = Instance(callee_manager_info, [model_tvar_type])
# instantiate with a proper model, Manager[MyModel], filling all Manager type vars in process
new_manager_info = self.new_typeinfo(self.class_name,
bases=[parent_manager_type])
new_manager_info.defn.type_vars = [model_tvar_defn]
new_manager_info.type_vars = [model_tvar_defn.name]
new_manager_info.set_line(self.call_expr)
# copy methods from passed_queryset_info with self type replaced
self_type = Instance(new_manager_info, [model_tvar_type])
for name, method_node in iter_all_custom_queryset_methods(passed_queryset_info):
self.add_method_from_signature(method_node,
name,
self_type,
new_manager_info.defn)
new_manager_sym = SymbolTableNode(self.semanal_api.current_symbol_kind(),
new_manager_info,
plugin_generated=True)
self.semanal_api.add_symbol_table_node(self.class_name, new_manager_sym)
# add mapping between generated manager and current one
runtime_manager_class_name = None
if 'class_name' in self.call_expr.arg_names:
class_name_arg = self.call_expr.args[self.call_expr.arg_names.index('class_name')]
if isinstance(class_name_arg, StrExpr):
runtime_manager_class_name = class_name_arg.value
new_manager_name = runtime_manager_class_name or (callee_manager_info.name + 'From' + queryset_class_name)
django_generated_manager_name = 'django.db.models.manager.' + new_manager_name
base_model_info.metadata.setdefault('managers', {})[django_generated_manager_name] = new_manager_info.fullname

View File

@@ -0,0 +1,236 @@
from abc import abstractmethod
from typing import Type, Optional
from django.db.models.base import Model
from django.db.models.fields.related import OneToOneField, ForeignKey
from django.db.models.fields.reverse_related import OneToOneRel, ManyToManyRel, ManyToOneRel
from mypy.checker import gen_unique_name
from mypy.nodes import TypeInfo, Var, SymbolTableNode, MDEF
from mypy.plugin import ClassDefContext
from mypy.semanal import dummy_context
from mypy.types import Instance, TypeOfAny, AnyType
from mypy.types import Type as MypyType
from django.db import models
from mypy_django_plugin.lib import helpers, fullnames
from mypy_django_plugin.transformers import fields
from mypy_django_plugin.transformers.fields import get_field_type
from mypy_django_plugin.transformers2 import new_helpers
class TransformModelClassCallback(helpers.ClassDefPluginCallback):
def get_real_manager_fullname(self, manager_fullname: str) -> str:
model_info = self.lookup_typeinfo_or_defer(fullnames.MODEL_CLASS_FULLNAME)
real_manager_fullname = model_info.metadata.get('managers', {}).get(manager_fullname, manager_fullname)
return real_manager_fullname
def modify_class_defn(self) -> None:
model_cls = self.django_context.get_model_class_by_fullname(self.class_defn.fullname)
if model_cls is None:
return None
return self.modify_model_class_defn(model_cls)
def add_new_model_attribute(self, name: str, typ: MypyType, force_replace: bool = False) -> None:
model_info = self.class_defn.info
if name in model_info.names and not force_replace:
raise ValueError('Attribute already exists on the model')
var = Var(name, type=typ)
var.info = model_info
var._fullname = self.semanal_api.qualified_name(name)
var.is_initialized_in_class = True
sym = SymbolTableNode(MDEF, var, plugin_generated=True)
error_context = None if force_replace else dummy_context()
added = self.semanal_api.add_symbol_table_node(name, sym, context=error_context)
assert added
def lookup_typeinfo_for_class_or_defer(self, klass: type, *,
reason_for_defer: Optional[str] = None) -> Optional[TypeInfo]:
manager_cls_fullname = helpers.get_class_fullname(klass)
return self.lookup_typeinfo_or_defer(manager_cls_fullname,
reason_for_defer=reason_for_defer)
@abstractmethod
def modify_model_class_defn(self, runtime_model_cls: Type[Model]) -> None:
raise NotImplementedError
class AddDefaultManagerCallback(TransformModelClassCallback):
def modify_model_class_defn(self, runtime_model_cls: Type[Model]) -> None:
if ('_default_manager' in self.class_defn.info.names
or runtime_model_cls._meta.default_manager is None):
return None
runtime_default_manager_class = runtime_model_cls._meta.default_manager.__class__
runtime_manager_cls_fullname = new_helpers.get_class_fullname(runtime_default_manager_class)
manager_cls_fullname = self.get_real_manager_fullname(runtime_manager_cls_fullname)
default_manager_info = self.lookup_typeinfo_or_defer(manager_cls_fullname)
if default_manager_info is None:
return
self.add_new_model_attribute('_default_manager',
Instance(default_manager_info, [Instance(self.class_defn.info, [])]))
class AddManagersCallback(TransformModelClassCallback):
def modify_model_class_defn(self, runtime_model_cls: Type[models.Model]) -> None:
for manager_name, manager in runtime_model_cls._meta.managers_map.items():
if manager_name in self.class_defn.info.names:
# already defined on the current model class, in file or at a previous iteration
continue
manager_info = self.lookup_typeinfo_for_class_or_defer(manager.__class__)
if manager_info is None:
continue
manager_type = Instance(manager_info, [Instance(self.class_defn.info, [])])
self.add_new_model_attribute(manager_name, manager_type)
class AddPrimaryKeyIfDoesNotExist(TransformModelClassCallback):
"""
Adds default primary key to models which does not define their own.
class User(models.Model):
name = models.TextField()
"""
def modify_model_class_defn(self, runtime_model_cls: Type[Model]) -> None:
auto_pk_field = runtime_model_cls._meta.auto_field
if auto_pk_field is None:
# defined explicitly
return None
auto_pk_field_name = auto_pk_field.attname
if auto_pk_field_name in self.class_defn.info.names:
# added on previous iteration
return None
auto_pk_field_info = self.lookup_typeinfo_for_class_or_defer(auto_pk_field.__class__)
if auto_pk_field_info is None:
return None
self.add_new_model_attribute(auto_pk_field_name,
fields.get_field_type(auto_pk_field_info, is_nullable=False))
class AddRelatedManagersCallback(TransformModelClassCallback):
def modify_model_class_defn(self, runtime_model_cls: Type[Model]) -> None:
for relation in self.django_context.get_model_relations(runtime_model_cls):
reverse_manager_name = relation.get_accessor_name()
if (reverse_manager_name is None
or reverse_manager_name in self.class_defn.info.names):
continue
related_model_cls = self.django_context.get_field_related_model_cls(relation)
if related_model_cls is None:
# could not find a referenced model (maybe invalid to= value, or GenericForeignKey)
continue
related_model_info = self.lookup_typeinfo_for_class_or_defer(related_model_cls)
if related_model_info is None:
continue
if isinstance(relation, OneToOneRel):
self.add_new_model_attribute(reverse_manager_name,
Instance(related_model_info, []))
elif isinstance(relation, (ManyToOneRel, ManyToManyRel)):
related_manager_info = self.lookup_typeinfo_or_defer(fullnames.RELATED_MANAGER_CLASS)
if related_manager_info is None:
if not self.defer_till_next_iteration(self.class_defn,
reason=f'{fullnames.RELATED_MANAGER_CLASS!r} is not available for lookup'):
raise TypeInfoNotFound(fullnames.RELATED_MANAGER_CLASS)
continue
# get type of default_manager for model
default_manager_fullname = helpers.get_class_fullname(related_model_cls._meta.default_manager.__class__)
reason_for_defer = (f'Trying to lookup default_manager {default_manager_fullname!r} '
f'of model {helpers.get_class_fullname(related_model_cls)!r}')
default_manager_info = self.lookup_typeinfo_or_defer(default_manager_fullname,
reason_for_defer=reason_for_defer)
if default_manager_info is None:
continue
default_manager_type = Instance(default_manager_info, [Instance(related_model_info, [])])
# related_model_cls._meta.default_manager.__class__
# # we're making a subclass of 'objects', need to have it defined
# if 'objects' not in related_model_info.names:
# if not self.defer_till_next_iteration(self.class_defn,
# reason=f"'objects' manager is not yet defined on {related_model_info.fullname!r}"):
# raise AttributeNotFound(self.class_defn.info, 'objects')
# continue
related_manager_type = Instance(related_manager_info,
[Instance(related_model_info, [])])
#
# objects_sym = related_model_info.names['objects']
# default_manager_type = objects_sym.type
# if default_manager_type is None:
# # dynamic base class, extract from django_context
# default_manager_cls = related_model_cls._meta.default_manager.__class__
# default_manager_info = self.lookup_typeinfo_for_class_or_defer(default_manager_cls)
# if default_manager_info is None:
# continue
# default_manager_type = Instance(default_manager_info, [Instance(related_model_info, [])])
if (not isinstance(default_manager_type, Instance)
or default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME):
# if not defined or trivial -> just return RelatedManager[Model]
self.add_new_model_attribute(reverse_manager_name, related_manager_type)
continue
# make anonymous class
name = gen_unique_name(related_model_cls.__name__ + '_' + 'RelatedManager',
self.semanal_api.current_symbol_table())
bases = [related_manager_type, default_manager_type]
new_manager_info = self.new_typeinfo(name, bases)
self.add_new_model_attribute(reverse_manager_name, Instance(new_manager_info, []))
class AddForeignPrimaryKeys(TransformModelClassCallback):
def modify_model_class_defn(self, runtime_model_cls: Type[Model]) -> None:
for field in runtime_model_cls._meta.get_fields():
if not isinstance(field, (OneToOneField, ForeignKey)):
continue
rel_pk_field_name = field.attname
if rel_pk_field_name in self.class_defn.info.names:
continue
related_model_cls = self.django_context.get_field_related_model_cls(field)
if related_model_cls is None:
field_sym = self.class_defn.info.get(field.name)
if field_sym is not None and field_sym.node is not None:
error_context = field_sym.node
else:
error_context = self.class_defn
self.semanal_api.fail(f'Cannot find model {field.related_model!r} '
f'referenced in field {field.name!r} ',
ctx=error_context)
self.add_new_model_attribute(rel_pk_field_name, AnyType(TypeOfAny.from_error))
continue
if related_model_cls._meta.abstract:
continue
rel_pk_field = self.django_context.get_primary_key_field(related_model_cls)
rel_pk_field_info = self.lookup_typeinfo_for_class_or_defer(rel_pk_field.__class__)
if rel_pk_field_info is None:
continue
field_type = get_field_type(rel_pk_field_info,
is_nullable=self.django_context.get_field_nullability(field))
self.add_new_model_attribute(rel_pk_field_name, field_type)
class ModelCallback(helpers.ClassDefPluginCallback):
def __call__(self, ctx: ClassDefContext) -> None:
callback_classes = [
AddManagersCallback,
AddPrimaryKeyIfDoesNotExist,
AddForeignPrimaryKeys,
AddDefaultManagerCallback,
AddRelatedManagersCallback,
]
for callback_cls in callback_classes:
callback = callback_cls(self.plugin)
callback.__call__(ctx)

View File

@@ -0,0 +1,26 @@
from typing import Union
from mypy.nodes import TypeInfo, MypyFile
class IncompleteDefnError(Exception):
pass
class TypeInfoNotFound(IncompleteDefnError):
def __init__(self, fullname: str) -> None:
super().__init__(f'It is final iteration and required type {fullname!r} is not ready yet.')
class AttributeNotFound(IncompleteDefnError):
def __init__(self, node: Union[TypeInfo, MypyFile], attrname: str) -> None:
super().__init__(f'Attribute {attrname!r} is not defined for the {node.fullname!r}.')
class NameNotFound(IncompleteDefnError):
def __init__(self, name: str) -> None:
super().__init__(f'Could not find {name!r} in the current activated namespaces')
def get_class_fullname(klass: type) -> str:
return klass.__module__ + '.' + klass.__qualname__

View File

@@ -36,7 +36,7 @@
pass
class Book(models.Model):
publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
owner = models.ForeignKey(db_column='model_id', to='auth.User', on_delete=models.CASCADE)
owner = models.ForeignKey(to='auth.User', on_delete=models.CASCADE)
- case: foreign_key_field_different_order_of_params
main: |
@@ -653,7 +653,7 @@
- case: related_manager_is_a_subclass_of_default_manager
main: |
from myapp.models import User
reveal_type(User().orders) # N: Revealed type is 'myapp.models.Order_RelatedManager'
reveal_type(User().orders) # N: Revealed type is 'myapp.models.User.Order_RelatedManager'
reveal_type(User().orders.get()) # N: Revealed type is 'myapp.models.Order*'
reveal_type(User().orders.manager_method()) # N: Revealed type is 'builtins.int'
installed_apps:
@@ -663,9 +663,12 @@
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class User(models.Model):
pass
class OrderManager(models.Manager):
class OrderManager(models.Manager['Order']):
def manager_method(self) -> int:
pass
class Order(models.Model):

View File

@@ -1,10 +1,133 @@
- case: from_queryset_with_base_manager
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.NewBaseManager[myapp.models.MyModel*]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.queryset_method) # N: Revealed type is 'def () -> builtins.str'
reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is 'builtins.str'
reveal_type(MyModel().objects2) # N: Revealed type is 'myapp.models.NewManager[myapp.models.MyModel*]'
reveal_type(MyModel().objects2.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects2.queryset_method) # N: Revealed type is 'def () -> builtins.str'
reveal_type(MyModel().objects2.queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from django.db.models.manager import BaseManager, Manager
from mypy_django_plugin.lib import generics
class ModelQuerySet(models.QuerySet):
def queryset_method(self) -> str:
return 'hello'
NewBaseManager = BaseManager.from_queryset(ModelQuerySet)
NewManager = Manager.from_queryset(ModelQuerySet)
generics.make_classes_generic(NewBaseManager, NewManager)
class MyModel(models.Model):
objects = NewBaseManager['MyModel']()
objects2 = NewManager['MyModel']()
- case: manager_without_generic_requires_annotation
main: |
from myapp.models import ModelQuerySet
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from django.db.models.manager import BaseManager, Manager
from mypy_django_plugin.lib import generics
class ModelQuerySet(models.QuerySet):
def queryset_method(self) -> str:
return 'hello'
NewBaseManager = BaseManager.from_queryset(ModelQuerySet)
NewManager = Manager.from_queryset(ModelQuerySet)
generics.make_classes_generic(NewBaseManager, NewManager)
class MyModel(models.Model):
objects = NewBaseManager() # E: Need type annotation for 'objects'
- case: from_queryset_with_custom_manager_as_base
main: |
from myapp.models import MyModel, NewManager
reveal_type(NewManager()) # N: Revealed type is 'myapp.models.NewManager[<nothing>]'
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.NewManager[myapp.models.MyModel]'
reveal_type(MyModel.objects.get()) # N: Revealed type is 'Any'
reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is 'builtins.int'
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
class ModelBaseManager(models.Manager):
def manager_only_method(self) -> int:
return 1
class ModelQuerySet(models.QuerySet):
def manager_and_queryset_method(self) -> str:
return 'hello'
NewManager = ModelBaseManager.from_queryset(ModelQuerySet)
generics.make_classes_generic(NewManager)
class MyModel(models.Model):
objects = NewManager['MyModel']()
- case: from_queryset_with_class_name_provided
main: |
from myapp.models import MyModel, NewManager
reveal_type(NewManager()) # N: Revealed type is 'myapp.models.NewManager[<nothing>]'
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.NewManager[myapp.models.MyModel]'
reveal_type(MyModel.objects.get()) # N: Revealed type is 'Any'
reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is 'builtins.int'
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
class ModelBaseManager(models.Manager):
def manager_only_method(self) -> int:
return 1
class ModelQuerySet(models.QuerySet):
def manager_and_queryset_method(self) -> str:
return 'hello'
NewManager = ModelBaseManager.from_queryset(ModelQuerySet, class_name='NewManager')
from mypy_django_plugin.lib import generics
generics.make_classes_generic(NewManager)
class MyModel(models.Model):
objects = NewManager['MyModel']()
- case: from_queryset_with_class_inheritance
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.NewManager[myapp.models.MyModel*]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
@@ -13,116 +136,24 @@
content: |
from django.db import models
from django.db.models.manager import BaseManager
class ModelQuerySet(models.QuerySet):
class BaseQuerySet(models.QuerySet):
def queryset_method(self) -> str:
return 'hello'
class ModelQuerySet(BaseQuerySet):
pass
NewManager = BaseManager.from_queryset(ModelQuerySet)
from mypy_django_plugin.lib import generics
generics.make_classes_generic(NewManager)
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_with_manager
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class ModelQuerySet(models.QuerySet):
def queryset_method(self) -> str:
return 'hello'
NewManager = models.Manager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_returns_intersection_of_manager_and_queryset
main: |
from myapp.models import MyModel, NewManager
reveal_type(NewManager()) # N: Revealed type is 'myapp.models.NewManager'
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel.objects.get()) # N: Revealed type is 'Any'
reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is 'builtins.int'
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class ModelBaseManager(models.Manager):
def manager_only_method(self) -> int:
return 1
class ModelQuerySet(models.QuerySet):
def manager_and_queryset_method(self) -> str:
return 'hello'
NewManager = ModelBaseManager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_with_class_name_provided
main: |
from myapp.models import MyModel, NewManager
reveal_type(NewManager()) # N: Revealed type is 'myapp.models.NewManager'
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel.objects.get()) # N: Revealed type is 'Any'
reveal_type(MyModel.objects.manager_only_method()) # N: Revealed type is 'builtins.int'
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class ModelBaseManager(models.Manager):
def manager_only_method(self) -> int:
return 1
class ModelQuerySet(models.QuerySet):
def manager_and_queryset_method(self) -> str:
return 'hello'
NewManager = ModelBaseManager.from_queryset(ModelQuerySet, class_name='NewManager')
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_with_class_inheritance
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from django.db.models.manager import BaseManager
class BaseQuerySet(models.QuerySet):
def queryset_method(self) -> str:
return 'hello'
class ModelQuerySet(BaseQuerySet):
pass
NewManager = BaseManager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()
objects = NewManager['MyModel']()
- case: from_queryset_with_manager_in_another_directory_and_imports
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.managers.NewManager[myapp.models.MyModel*]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.queryset_method) # N: Revealed type is 'def (param: Union[builtins.str, None] =) -> Union[builtins.str, None]'
reveal_type(MyModel().objects.queryset_method('str')) # N: Revealed type is 'Union[builtins.str, None]'
@@ -135,8 +166,11 @@
from django.db import models
from myapp.managers import NewManager
from mypy_django_plugin.lib import generics
generics.make_classes_generic(NewManager)
class MyModel(models.Model):
objects = NewManager()
objects = NewManager['MyModel']()
- path: myapp/managers.py
content: |
from typing import Optional
@@ -152,7 +186,7 @@
disable_cache: true
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.managers.NewManager[myapp.models.MyModel*]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.base_queryset_method) # N: Revealed type is 'def (param: Union[builtins.int, builtins.str]) -> <nothing>'
reveal_type(MyModel().objects.base_queryset_method(2)) # N: Revealed type is '<nothing>'
@@ -164,8 +198,12 @@
content: |
from django.db import models
from myapp.managers import NewManager
from mypy_django_plugin.lib import generics
generics.make_classes_generic(NewManager)
class MyModel(models.Model):
objects = NewManager()
objects = NewManager['MyModel']()
- path: myapp/managers.py
content: |
from django.db import models
@@ -192,7 +230,7 @@
reveal_type(Permission().another_models) # N: Revealed type is 'django.db.models.manager.RelatedManager[myapp.models.AnotherModelInProjectWithContribAuthM2M]'
from myapp.managers import NewManager
reveal_type(NewManager()) # N: Revealed type is 'myapp.managers.NewManager'
reveal_type(NewManager()) # N: Revealed type is 'myapp.managers.NewManager[<nothing>]'
reveal_type(NewManager().base_queryset_method) # N: Revealed type is 'def (param: builtins.dict[builtins.str, Union[builtins.int, builtins.str]]) -> Union[builtins.int, builtins.str]'
from myapp.models import MyModel
@@ -210,8 +248,11 @@
from myapp.managers import NewManager
from django.contrib.auth.models import Permission
from mypy_django_plugin.lib import generics
generics.make_classes_generic(NewManager)
class MyModel(models.Model):
objects = NewManager()
objects = NewManager['MyModel']()
class AnotherModelInProjectWithContribAuthM2M(models.Model):
permissions = models.ForeignKey(

View File

@@ -220,6 +220,9 @@
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.QuerySet)
class TransactionQuerySet(models.QuerySet['Transaction']):
pass
class Transaction(models.Model):

View File

@@ -30,6 +30,7 @@
class Child(Parent):
pass
- case: test_model_objects_attribute_present_in_case_of_model_cls_passed_as_generic_parameter
main: |
from myapp.models import Base, MyModel
@@ -55,23 +56,27 @@
def method(self) -> None:
reveal_type(self.model_cls._default_manager) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.MyModel]'
- case: if_custom_manager_defined_it_is_set_to_default_manager
main: |
from myapp.models import MyModel
reveal_type(MyModel._default_manager) # N: Revealed type is 'myapp.models.CustomManager[myapp.models.MyModel]'
reveal_type(MyModel._default_manager.get()) # N: Revealed type is 'myapp.models.MyModel*'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing import TypeVar
from django.db import models
_T = TypeVar('_T', bound=models.Model)
class CustomManager(models.Manager[_T]):
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class CustomManager(models.Manager['MyModel']):
pass
class MyModel(models.Model):
manager = CustomManager['MyModel']()
manager = CustomManager()
- case: if_default_manager_name_is_passed_set_default_manager_to_it
main: |
@@ -83,40 +88,54 @@
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing import TypeVar
from django.db import models
_T = TypeVar('_T', bound=models.Model)
class Manager1(models.Manager[_T]):
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class Manager1(models.Manager['MyModel']):
pass
class Manager2(models.Manager[_T]):
class Manager2(models.Manager['MyModel']):
pass
class MyModel(models.Model):
class Meta:
default_manager_name = 'm2'
m1 = Manager1['MyModel']()
m2 = Manager2['MyModel']()
m1 = Manager1()
m2 = Manager2()
- case: test_leave_as_is_if_objects_is_set_and_fill_typevars_with_outer_class
- case: manager_requires_type_annotation_to_be_set_if_generic_is_not_specified
main: |
from myapp.models import MyUser
reveal_type(MyUser.objects) # N: Revealed type is 'myapp.models.UserManager[myapp.models.MyUser]'
reveal_type(MyUser.objects.get()) # N: Revealed type is 'myapp.models.MyUser*'
reveal_type(MyUser.objects.get_or_404()) # N: Revealed type is 'myapp.models.MyUser'
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.MyManager'
reveal_type(MyModel.objects.get()) # N: Revealed type is 'Any'
reveal_type(MyModel.objects2) # N: Revealed type is 'myapp.models.MyGenericManager[Any]'
reveal_type(MyModel.objects2.get()) # N: Revealed type is 'Any'
reveal_type(MyModel.objects3) # N: Revealed type is 'myapp.models.MyGenericManager[myapp.models.MyModel]'
reveal_type(MyModel.objects3.get()) # N: Revealed type is 'myapp.models.MyModel*'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from typing import TypeVar
from django.db import models
class UserManager(models.Manager['MyUser']):
def get_or_404(self) -> 'MyUser':
pass
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class MyManager(models.Manager):
pass
_T = TypeVar('_T', bound=models.Model)
class MyGenericManager(models.Manager[_T]):
pass
class MyModel(models.Model):
objects = MyManager()
objects2 = MyGenericManager() # E: Need type annotation for 'objects2'
objects3: 'MyGenericManager[MyModel]' = MyGenericManager()
class MyUser(models.Model):
objects = UserManager()
- case: model_imported_from_different_file
main: |
@@ -139,13 +158,14 @@
class Inventory(models.Model):
pass
- case: managers_that_defined_on_other_models_do_not_influence
main: |
from myapp.models import AbstractPerson, Book
reveal_type(AbstractPerson.abstract_persons) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.AbstractPerson]'
reveal_type(Book.published_objects) # N: Revealed type is 'myapp.models.PublishedBookManager[myapp.models.Book]'
reveal_type(Book.published_objects) # N: Revealed type is 'myapp.models.PublishedBookManager'
Book.published_objects.create(title='hello')
reveal_type(Book.annotated_objects) # N: Revealed type is 'myapp.models.AnnotatedBookManager[myapp.models.Book]'
reveal_type(Book.annotated_objects) # N: Revealed type is 'myapp.models.AnnotatedBookManager'
Book.annotated_objects.create(title='hello')
installed_apps:
- myapp
@@ -155,6 +175,9 @@
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class AbstractPerson(models.Model):
abstract_persons = models.Manager['AbstractPerson']()
class PublishedBookManager(models.Manager['Book']):
@@ -166,7 +189,8 @@
published_objects = PublishedBookManager()
annotated_objects = AnnotatedBookManager()
- case: managers_inherited_from_abstract_classes_multiple_inheritance
- case: managers_inherited_from_abstract_classes_multiple_inheritance_do_not_warn_on_liskov
main: |
installed_apps:
- myapp
@@ -175,6 +199,9 @@
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class CustomManager1(models.Manager['AbstractBase1']):
pass
class AbstractBase1(models.Model):
@@ -193,6 +220,7 @@
class Child(AbstractBase1, AbstractBase2):
pass
- case: model_has_a_manager_of_the_same_type
main: |
from myapp.models import UnrelatedModel, MyModel
@@ -208,59 +236,21 @@
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class UnrelatedModel(models.Model):
objects = models.Manager['UnrelatedModel']()
class MyModel(models.Model):
pass
- case: manager_without_annotation_of_the_model_gets_it_from_outer_one
main: |
from myapp.models import UnrelatedModel2, MyModel2
reveal_type(UnrelatedModel2.objects) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.UnrelatedModel2]'
reveal_type(UnrelatedModel2.objects.first()) # N: Revealed type is 'Union[myapp.models.UnrelatedModel2*, None]'
reveal_type(MyModel2.objects) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.MyModel2]'
reveal_type(MyModel2.objects.first()) # N: Revealed type is 'Union[myapp.models.MyModel2*, None]'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class UnrelatedModel2(models.Model):
objects = models.Manager()
class MyModel2(models.Model):
pass
- case: inherited_manager_has_the_proper_type_of_model
main: |
from myapp.models import ParentOfMyModel3, MyModel3
reveal_type(ParentOfMyModel3.objects) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.ParentOfMyModel3]'
reveal_type(ParentOfMyModel3.objects.first()) # N: Revealed type is 'Union[myapp.models.ParentOfMyModel3*, None]'
reveal_type(MyModel3.objects) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.MyModel3]'
reveal_type(MyModel3.objects.first()) # N: Revealed type is 'Union[myapp.models.MyModel3*, None]'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class ParentOfMyModel3(models.Model):
objects = models.Manager()
class MyModel3(ParentOfMyModel3):
pass
- case: inheritance_with_explicit_type_on_child_manager
main: |
from myapp.models import ParentOfMyModel4, MyModel4
reveal_type(ParentOfMyModel4.objects) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.ParentOfMyModel4]'
reveal_type(ParentOfMyModel4.objects.first()) # N: Revealed type is 'Union[myapp.models.ParentOfMyModel4*, None]'
reveal_type(ParentOfMyModel4.objects) # N: Revealed type is 'django.db.models.manager.Manager[Any]'
reveal_type(ParentOfMyModel4.objects.first()) # N: Revealed type is 'Union[Any, None]'
reveal_type(MyModel4.objects) # N: Revealed type is 'django.db.models.manager.Manager[myapp.models.MyModel4]'
reveal_type(MyModel4.objects.first()) # N: Revealed type is 'Union[myapp.models.MyModel4*, None]'
@@ -271,76 +261,20 @@
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class ParentOfMyModel4(models.Model):
objects = models.Manager()
objects = models.Manager() # E: Need type annotation for 'objects'
class MyModel4(ParentOfMyModel4):
objects = models.Manager['MyModel4']()
# TODO: make it work someday
#- case: inheritance_of_two_models_with_custom_objects_manager
# main: |
# from myapp.models import MyBaseUser, MyUser
# reveal_type(MyBaseUser.objects) # N: Revealed type is 'myapp.models.MyBaseManager[myapp.models.MyBaseUser]'
# reveal_type(MyBaseUser.objects.get()) # N: Revealed type is 'myapp.models.MyBaseUser'
#
# reveal_type(MyUser.objects) # N: Revealed type is 'myapp.models.MyManager[myapp.models.MyUser]'
# reveal_type(MyUser.objects.get()) # N: Revealed type is 'myapp.models.MyUser'
# installed_apps:
# - myapp
# files:
# - path: myapp/__init__.py
# - path: myapp/models.py
# content: |
# from django.db import models
#
# class MyBaseManager(models.Manager):
# pass
# class MyBaseUser(models.Model):
# objects = MyBaseManager()
#
# class MyManager(models.Manager):
# pass
# class MyUser(MyBaseUser):
# objects = MyManager()
- case: custom_manager_returns_proper_model_types
main: |
from myapp.models import User
reveal_type(User.objects) # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]'
reveal_type(User.objects.select_related()) # N: Revealed type is 'myapp.models.User_MyManager[myapp.models.User]'
reveal_type(User.objects.get()) # N: Revealed type is 'myapp.models.User*'
reveal_type(User.objects.get_instance()) # N: Revealed type is 'builtins.int'
reveal_type(User.objects.get_instance_untyped('hello')) # N: Revealed type is 'Any'
from myapp.models import ChildUser
reveal_type(ChildUser.objects) # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]'
reveal_type(ChildUser.objects.select_related()) # N: Revealed type is 'myapp.models.ChildUser_MyManager[myapp.models.ChildUser]'
reveal_type(ChildUser.objects.get()) # N: Revealed type is 'myapp.models.ChildUser*'
reveal_type(ChildUser.objects.get_instance()) # N: Revealed type is 'builtins.int'
reveal_type(ChildUser.objects.get_instance_untyped('hello')) # N: Revealed type is 'Any'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
class MyManager(models.Manager):
def get_instance(self) -> int:
pass
def get_instance_untyped(self, name):
pass
class User(models.Model):
objects = MyManager()
class ChildUser(models.Model):
objects = MyManager()
- case: manager_defined_in_the_nested_class
- case: manager_defined_as_a_nested_class
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.MyModel_MyManager[myapp.models.MyModel]'
reveal_type(MyModel.objects) # N: Revealed type is 'myapp.models.MyModel.MyManager'
reveal_type(MyModel.objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel.objects.mymethod()) # N: Revealed type is 'builtins.int'
installed_apps:
@@ -350,8 +284,11 @@
- path: myapp/models.py
content: |
from django.db import models
from mypy_django_plugin.lib import generics
generics.make_classes_generic(models.Manager)
class MyModel(models.Model):
class MyManager(models.Manager):
class MyManager(models.Manager['MyModel']):
def mymethod(self) -> int:
pass
objects = MyManager()