Set custom queryset methods as manager attrs instead of method copies (#820)

Instead of copying methods over from a QuerySet passed to a basemanager
when invoking '<BaseManager>.from_queryset', any QuerySet methods are
declared as attributes on the manager.

This allows us to properly lookup any QuerySet method types via a
'get_attribute_hook' and will thus remove disorienting phantom errors
occuring from mypy trying to resolve types only existing in the module
where the _original_ (and real) queryset method was declared.
This commit is contained in:
Petter Friberg
2022-01-16 10:14:33 +01:00
committed by GitHub
parent 1da693ebff
commit 99f28387fb
6 changed files with 323 additions and 48 deletions

View File

@@ -34,7 +34,7 @@ from mypy.plugin import (
MethodContext, MethodContext,
SemanticAnalyzerPluginInterface, SemanticAnalyzerPluginInterface,
) )
from mypy.plugins.common import add_method from mypy.plugins.common import add_method_to_class
from mypy.semanal import SemanticAnalyzer from mypy.semanal import SemanticAnalyzer
from mypy.types import AnyType, CallableType, Instance, NoneTyp, TupleType from mypy.types import AnyType, CallableType, Instance, NoneTyp, TupleType
from mypy.types import Type as MypyType from mypy.types import Type as MypyType
@@ -383,7 +383,9 @@ def copy_method_to_another_class(
return return
arguments, return_type = build_unannotated_method_args(method_node) arguments, return_type = build_unannotated_method_args(method_node)
add_method(ctx, new_method_name, args=arguments, return_type=return_type, self_type=self_type) add_method_to_class(
semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
)
return return
method_type = method_node.type method_type = method_node.type
@@ -421,7 +423,9 @@ def copy_method_to_another_class(
argument.set_line(original_argument) argument.set_line(original_argument)
arguments.append(argument) arguments.append(argument)
add_method(ctx, new_method_name, args=arguments, return_type=return_type, self_type=self_type) add_method_to_class(
semanal_api, ctx.cls, new_method_name, args=arguments, return_type=return_type, self_type=self_type
)
def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None: def add_new_manager_base(api: SemanticAnalyzerPluginInterface, fullname: str) -> None:

View File

@@ -24,7 +24,10 @@ import mypy_django_plugin.transformers.orm_lookups
from mypy_django_plugin.django.context import DjangoContext from mypy_django_plugin.django.context import DjangoContext
from mypy_django_plugin.lib import fullnames, helpers from mypy_django_plugin.lib import fullnames, helpers
from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings
from mypy_django_plugin.transformers.managers import create_new_manager_class_from_from_queryset_method from mypy_django_plugin.transformers.managers import (
create_new_manager_class_from_from_queryset_method,
resolve_manager_method,
)
from mypy_django_plugin.transformers.models import ( from mypy_django_plugin.transformers.models import (
handle_annotated_type, handle_annotated_type,
process_model_class, process_model_class,
@@ -302,6 +305,7 @@ class NewSemanalDjangoPlugin(Plugin):
mypy_django_plugin.transformers.orm_lookups.typecheck_queryset_filter, mypy_django_plugin.transformers.orm_lookups.typecheck_queryset_filter,
django_context=self.django_context, django_context=self.django_context,
) )
return None return None
def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]: def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
@@ -330,6 +334,13 @@ class NewSemanalDjangoPlugin(Plugin):
return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context) return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context)
if info and info.has_base(fullnames.ABSTRACT_USER_MODEL_FULLNAME) and attr_name in ("is_staff", "is_active"): if info and info.has_base(fullnames.ABSTRACT_USER_MODEL_FULLNAME) and attr_name in ("is_staff", "is_active"):
return partial(set_auth_user_model_boolean_fields, django_context=self.django_context) return partial(set_auth_user_model_boolean_fields, django_context=self.django_context)
if (
info
and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
and class_name in self._get_current_manager_bases()
):
return partial(resolve_manager_method, django_context=self.django_context)
return None return None
def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeTypeContext], MypyType]]: def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeTypeContext], MypyType]]:

View File

@@ -1,14 +1,137 @@
from mypy.checker import fill_typevars from typing import Optional, Union
from mypy.nodes import GDEF, Decorator, FuncDef, MemberExpr, NameExpr, RefExpr, StrExpr, SymbolTableNode, TypeInfo
from mypy.plugin import ClassDefContext, DynamicClassDefContext
from mypy.types import CallableType, Instance, TypeVarType, UnboundType, get_proper_type
from mypy.checker import TypeChecker, fill_typevars
from mypy.nodes import (
GDEF,
Decorator,
FuncBase,
FuncDef,
MemberExpr,
NameExpr,
OverloadedFuncDef,
RefExpr,
StrExpr,
SymbolTableNode,
TypeInfo,
Var,
)
from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
from mypy.types import AnyType, CallableType, Instance, ProperType
from mypy.types import Type as MypyType
from mypy.types import TypeOfAny, TypeVarType, UnboundType, get_proper_type
from mypy_django_plugin.django.context import DjangoContext
from mypy_django_plugin.lib import fullnames, helpers from mypy_django_plugin.lib import fullnames, helpers
def get_method_type_from_dynamic_manager(
api: TypeChecker, method_name: str, manager_type_info: TypeInfo
) -> Optional[ProperType]:
"""
Attempt to resolve a method on a manager that was built from '.from_queryset'
"""
if (
"django" not in manager_type_info.metadata
or "from_queryset_manager" not in manager_type_info.metadata["django"]
):
# Manager isn't dynamically added
return None
queryset_fullname = manager_type_info.metadata["django"]["from_queryset_manager"]
assert isinstance(queryset_fullname, str)
queryset_info = helpers.lookup_fully_qualified_typeinfo(api, queryset_fullname)
assert queryset_info is not None
def get_funcdef_type(definition: Union[FuncBase, Decorator, None]) -> Optional[ProperType]:
# TODO: Handle @overload?
if isinstance(definition, FuncBase) and not isinstance(definition, OverloadedFuncDef):
return definition.type
elif isinstance(definition, Decorator):
return definition.func.type
return None
method_type = get_funcdef_type(queryset_info.get_method(method_name))
if method_type is None:
return None
assert isinstance(method_type, CallableType)
# Drop any 'self' argument as our manager is already initialized
return method_type.copy_modified(
arg_types=method_type.arg_types[1:],
arg_kinds=method_type.arg_kinds[1:],
arg_names=method_type.arg_names[1:],
)
def get_method_type_from_reverse_manager(
api: TypeChecker, method_name: str, manager_type_info: TypeInfo
) -> Optional[ProperType]:
"""
Attempts to resolve a reverse manager's method via the '_default_manager' manager on the related model
From Django docs:
"By default the RelatedManager used for reverse relations is a subclass of the default manager for that model."
Ref: https://docs.djangoproject.com/en/dev/topics/db/queries/#using-a-custom-reverse-manager
"""
is_reverse_manager = (
"django" in manager_type_info.metadata and "related_manager_to_model" in manager_type_info.metadata["django"]
)
if not is_reverse_manager:
return None
related_model_fullname = manager_type_info.metadata["django"]["related_manager_to_model"]
assert isinstance(related_model_fullname, str)
model_info = helpers.lookup_fully_qualified_typeinfo(api, related_model_fullname)
if model_info is None:
return None
# We should _always_ have a '_default_manager' on a model
assert "_default_manager" in model_info.names
assert isinstance(model_info.names["_default_manager"].node, Var)
manager_instance = model_info.names["_default_manager"].node.type
return (
get_method_type_from_dynamic_manager(api, method_name, manager_instance.type)
# TODO: Can we assert on None and Instance?
if manager_instance is not None and isinstance(manager_instance, Instance)
else None
)
def resolve_manager_method(ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
"""
A 'get_attribute_hook' that is intended to be invoked whenever the TypeChecker encounters
an attribute on a class that has 'django.db.models.BaseManager' as a base.
"""
api = helpers.get_typechecker_api(ctx)
# Skip (method) type that is currently something other than Any
if not isinstance(ctx.default_attr_type, AnyType):
return ctx.default_attr_type
# (Current state is:) We wouldn't end up here when looking up a method from a custom _manager_.
# That's why we only attempt to lookup the method for either a dynamically added or reverse manager.
assert isinstance(ctx.context, MemberExpr)
method_name = ctx.context.name
manager_instance = ctx.type
assert isinstance(manager_instance, Instance)
method_type = get_method_type_from_dynamic_manager(
api, method_name, manager_instance.type
) or get_method_type_from_reverse_manager(api, method_name, manager_instance.type)
return method_type if method_type is not None else ctx.default_attr_type
def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None: def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None:
"""
Insert a new manager class node for a: '<Name> = <Manager>.from_queryset(<QuerySet>)'.
When the assignment expression lives at module level.
"""
semanal_api = helpers.get_semanal_api(ctx) semanal_api = helpers.get_semanal_api(ctx)
# Don't redeclare the manager class if we've already defined it.
manager_node = semanal_api.lookup_current_scope(ctx.name)
if manager_node and isinstance(manager_node.node, TypeInfo):
# This is just a deferral run where our work is already finished
return
callee = ctx.call.callee callee = ctx.call.callee
assert isinstance(callee, MemberExpr) assert isinstance(callee, MemberExpr)
assert isinstance(callee.expr, RefExpr) assert isinstance(callee.expr, RefExpr)
@@ -54,9 +177,11 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
new_manager_info.defn.type_vars = base_manager_info.defn.type_vars new_manager_info.defn.type_vars = base_manager_info.defn.type_vars
new_manager_info.defn.line = ctx.call.line new_manager_info.defn.line = ctx.call.line
new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type() new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type()
# Stash the queryset fullname which was passed to .from_queryset
current_module = semanal_api.cur_mod_node # So that our 'resolve_manager_method' attribute hook can fetch the method from that QuerySet class
current_module.names[ctx.name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) new_manager_info.metadata.setdefault("django", {})
new_manager_info.metadata["django"].setdefault("from_queryset_manager", {})
new_manager_info.metadata["django"]["from_queryset_manager"] = derived_queryset_fullname
if len(ctx.call.args) > 1: if len(ctx.call.args) > 1:
expr = ctx.call.args[1] expr = ctx.call.args[1]
@@ -66,8 +191,7 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
custom_manager_generated_name = base_manager_info.name + "From" + derived_queryset_info.name 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]) 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.setdefault("from_queryset_managers", {})
base_manager_info.metadata["from_queryset_managers"] = {}
base_manager_info.metadata["from_queryset_managers"][custom_manager_generated_fullname] = new_manager_info.fullname base_manager_info.metadata["from_queryset_managers"][custom_manager_generated_fullname] = new_manager_info.fullname
# So that the plugin will reparameterize the manager when it is constructed inside of a Model definition # So that the plugin will reparameterize the manager when it is constructed inside of a Model definition
@@ -76,13 +200,10 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
class_def_context = ClassDefContext(cls=new_manager_info.defn, reason=ctx.call, api=semanal_api) class_def_context = ClassDefContext(cls=new_manager_info.defn, reason=ctx.call, api=semanal_api)
self_type = fill_typevars(new_manager_info) self_type = fill_typevars(new_manager_info)
assert isinstance(self_type, Instance) assert isinstance(self_type, Instance)
queryset_method_names = []
# we need to copy all methods in MRO before django.db.models.query.QuerySet # We collect and mark up all methods before django.db.models.query.QuerySet as class members
for class_mro_info in derived_queryset_info.mro: for class_mro_info in derived_queryset_info.mro:
if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME: if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME:
for name, sym in class_mro_info.names.items():
queryset_method_names.append(name)
break break
for name, sym in class_mro_info.names.items(): for name, sym in class_mro_info.names.items():
if isinstance(sym.node, FuncDef): if isinstance(sym.node, FuncDef):
@@ -91,10 +212,17 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
func_node = sym.node.func func_node = sym.node.func
else: else:
continue continue
helpers.copy_method_to_another_class(
class_def_context, self_type, new_method_name=name, method_node=func_node
)
# Insert the queryset method name as a class member. Note that the type of
# the method is set as Any. Figuring out the type is the job of the
# 'resolve_manager_method' attribute hook, which comes later.
#
# class BaseManagerFromMyQuerySet(BaseManager):
# queryset_method: Any = ...
#
helpers.add_new_sym_for_info(new_manager_info, name=name, sym_type=AnyType(TypeOfAny.special_form))
# we need to copy all methods in MRO before django.db.models.query.QuerySet
# Gather names of all BaseManager methods # Gather names of all BaseManager methods
manager_method_names = [] manager_method_names = []
for manager_mro_info in new_manager_info.mro: for manager_mro_info in new_manager_info.mro:
@@ -150,3 +278,6 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
return_type=return_type, return_type=return_type,
original_module_name=class_mro_info.module_name, original_module_name=class_mro_info.module_name,
) )
# Insert the new manager (dynamic) class
assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True))

View File

@@ -68,6 +68,20 @@ class ModelClassInitializer:
return return
self.run_with_model_cls(model_cls) self.run_with_model_cls(model_cls)
def get_generated_manager_mappings(self, base_manager_fullname: str) -> Dict[str, str]:
base_manager_info = self.lookup_typeinfo(base_manager_fullname)
if base_manager_info is None or "from_queryset_managers" not in base_manager_info.metadata:
return {}
return base_manager_info.metadata["from_queryset_managers"]
def get_generated_manager_info(self, manager_fullname: str, base_manager_fullname: str) -> Optional[TypeInfo]:
generated_managers = self.get_generated_manager_mappings(base_manager_fullname)
real_manager_fullname = generated_managers.get(manager_fullname)
if real_manager_fullname:
return self.lookup_typeinfo(real_manager_fullname)
# Not a generated manager
return None
def run_with_model_cls(self, model_cls): def run_with_model_cls(self, model_cls):
raise NotImplementedError("Implement this in subclasses") raise NotImplementedError("Implement this in subclasses")
@@ -179,12 +193,6 @@ class AddManagers(ModelClassInitializer):
def is_any_parametrized_manager(self, typ: Instance) -> bool: def is_any_parametrized_manager(self, typ: Instance) -> bool:
return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType) return typ.type.fullname in fullnames.MANAGER_CLASSES and isinstance(typ.args[0], AnyType)
def get_generated_manager_mappings(self, base_manager_fullname: str) -> Dict[str, str]:
base_manager_info = self.lookup_typeinfo(base_manager_fullname)
if base_manager_info is None or "from_queryset_managers" not in base_manager_info.metadata:
return {}
return base_manager_info.metadata["from_queryset_managers"]
def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance: def create_new_model_parametrized_manager(self, name: str, base_manager_info: TypeInfo) -> Instance:
bases = [] bases = []
for original_base in base_manager_info.bases: for original_base in base_manager_info.bases:
@@ -230,23 +238,25 @@ class AddManagers(ModelClassInitializer):
if not self.api.final_iteration: if not self.api.final_iteration:
raise exc raise exc
else: else:
# On final round, see if we can find info for a generated (dynamic class) manager
base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0]) base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0])
generated_managers = self.get_generated_manager_mappings(base_manager_fullname) manager_info = self.get_generated_manager_info(manager_fullname, base_manager_fullname)
if manager_fullname not in generated_managers:
# not a generated manager, continue with the loop
continue
real_manager_fullname = generated_managers[manager_fullname]
manager_info = self.lookup_typeinfo(real_manager_fullname)
if manager_info is None: if manager_info is None:
continue continue
manager_class_name = real_manager_fullname.rsplit(".", maxsplit=1)[1] _, manager_class_name = manager_info.fullname.rsplit(".", maxsplit=1)
if manager_name not in self.model_classdef.info.names: if manager_name not in self.model_classdef.info.names:
manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])])
self.add_new_node_to_model_class(manager_name, manager_type) self.add_new_node_to_model_class(manager_name, manager_type)
else: else:
# creates new MODELNAME_MANAGERCLASSNAME class that represents manager parametrized with current model # Ending up here could for instance be due to having a custom _Manager_
if not self.has_any_parametrized_manager_as_base(manager_info): # that is not built from a custom QuerySet. Another example is a
# related manager.
# Don't interfere with dynamically generated manager classes
is_dynamically_generated = "django" in manager_info.metadata and manager_info.metadata["django"].get(
"from_queryset_manager"
)
if not self.has_any_parametrized_manager_as_base(manager_info) or is_dynamically_generated:
continue continue
custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name custom_model_manager_name = manager.model.__name__ + "_" + manager_class_name
@@ -262,12 +272,27 @@ class AddManagers(ModelClassInitializer):
class AddDefaultManagerAttribute(ModelClassInitializer): class AddDefaultManagerAttribute(ModelClassInitializer):
def run_with_model_cls(self, model_cls: Type[Model]) -> None: def run_with_model_cls(self, model_cls: Type[Model]) -> None:
# add _default_manager if "_default_manager" in self.model_classdef.info.names:
if "_default_manager" not in self.model_classdef.info.names: return None
default_manager_fullname = helpers.get_class_fullname(model_cls._meta.default_manager.__class__)
default_manager_cls = model_cls._meta.default_manager.__class__
default_manager_fullname = helpers.get_class_fullname(default_manager_cls)
try:
default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(default_manager_fullname) default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(default_manager_fullname)
default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])]) except helpers.IncompleteDefnException as exc:
self.add_new_node_to_model_class("_default_manager", default_manager) if not self.api.final_iteration:
raise exc
else:
base_manager_fullname = helpers.get_class_fullname(default_manager_cls.__bases__[0])
generated_manager_info = self.get_generated_manager_info(
default_manager_fullname, base_manager_fullname
)
if generated_manager_info is None:
return
default_manager_info = generated_manager_info
default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])])
self.add_new_node_to_model_class("_default_manager", default_manager)
class AddRelatedManagers(ModelClassInitializer): class AddRelatedManagers(ModelClassInitializer):
@@ -300,6 +325,8 @@ class AddRelatedManagers(ModelClassInitializer):
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error( related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
fullnames.RELATED_MANAGER_CLASS fullnames.RELATED_MANAGER_CLASS
) # noqa: E501 ) # noqa: E501
# TODO: Use default manager instead of 'objects'
# See: https://docs.djangoproject.com/en/dev/topics/db/queries/#using-a-custom-reverse-manager
objects = related_model_info.get("objects") objects = related_model_info.get("objects")
if not objects: if not objects:
raise helpers.IncompleteDefnException() raise helpers.IncompleteDefnException()
@@ -325,15 +352,10 @@ class AddRelatedManagers(ModelClassInitializer):
name = model_cls.__name__ + "_" + related_model_cls.__name__ + "_" + "RelatedManager" name = model_cls.__name__ + "_" + related_model_cls.__name__ + "_" + "RelatedManager"
bases = [parametrized_related_manager_type, default_manager_type] bases = [parametrized_related_manager_type, default_manager_type]
new_related_manager_info = self.add_new_class_for_current_module(name, bases) new_related_manager_info = self.add_new_class_for_current_module(name, bases)
new_related_manager_info.metadata["django"] = {"related_manager_to_model": related_model_info.fullname}
self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, [])) self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, []))
def get_generated_manager_mappings(self, base_manager_fullname: str) -> Dict[str, str]:
base_manager_info = self.lookup_typeinfo(base_manager_fullname)
if base_manager_info is None or "from_queryset_managers" not in base_manager_info.metadata:
return {}
return base_manager_info.metadata["from_queryset_managers"]
def try_generate_related_manager( def try_generate_related_manager(
self, related_model_cls: Type[Model], related_model_info: TypeInfo self, related_model_cls: Type[Model], related_model_info: TypeInfo
) -> Optional[Instance]: ) -> Optional[Instance]:

View File

@@ -684,6 +684,7 @@
reveal_type(User().orders) # N: Revealed type is "myapp.models.User_Order_RelatedManager1" reveal_type(User().orders) # N: Revealed type is "myapp.models.User_Order_RelatedManager1"
reveal_type(User().orders.get()) # N: Revealed type is "myapp.models.Order*" reveal_type(User().orders.get()) # N: Revealed type is "myapp.models.Order*"
reveal_type(User().orders.manager_method()) # N: Revealed type is "builtins.int" reveal_type(User().orders.manager_method()) # N: Revealed type is "builtins.int"
reveal_type(Product.objects.queryset_method()) # N: Revealed type is "builtins.int"
reveal_type(Order().products) # N: Revealed type is "myapp.models.Order_Product_RelatedManager1" reveal_type(Order().products) # N: Revealed type is "myapp.models.Order_Product_RelatedManager1"
reveal_type(Order().products.get()) # N: Revealed type is "myapp.models.Product*" reveal_type(Order().products.get()) # N: Revealed type is "myapp.models.Product*"
reveal_type(Order().products.queryset_method()) # N: Revealed type is "builtins.int" reveal_type(Order().products.queryset_method()) # N: Revealed type is "builtins.int"

View File

@@ -22,6 +22,55 @@
class MyModel(models.Model): class MyModel(models.Model):
objects = NewManager() objects = NewManager()
- case: from_queryset_queryset_imported_from_other_module
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.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 "myapp.querysets.ModelQuerySet"
reveal_type(MyModel.objects.queryset_method_2()) # N: Revealed type is "typing.Iterable[myapp.querysets.Custom]"
reveal_type(MyModel.objects.queryset_method_3()) # N: Revealed type is "builtins.str"
reveal_type(MyModel.objects.queryset_method_4([])) # N: Revealed type is "None"
reveal_type(MyModel.objects.filter(id=1).queryset_method()) # N: Revealed type is "myapp.querysets.ModelQuerySet"
reveal_type(MyModel.objects.filter(id=1)) # N: Revealed type is "myapp.querysets.ModelQuerySet[myapp.models.MyModel*]"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/querysets.py
content: |
from typing import TYPE_CHECKING, Iterable, Sequence
from django.db import models
if TYPE_CHECKING:
from .models import MyModel
class Custom:
...
class ModelQuerySet(models.QuerySet["MyModel"]):
def queryset_method(self) -> "ModelQuerySet":
return self.filter()
def queryset_method_2(self) -> Iterable[Custom]:
return []
def queryset_method_3(self) -> str:
return 'hello'
def queryset_method_4(self, arg: Sequence) -> None:
return None
- path: myapp/models.py
content: |
from django.db import models
from django.db.models.manager import BaseManager
from .querysets import ModelQuerySet
NewManager = BaseManager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_with_manager - case: from_queryset_with_manager
main: | main: |
from myapp.models import MyModel from myapp.models import MyModel
@@ -48,7 +97,7 @@
main: | main: |
from myapp.models import MyModel, NewManager from myapp.models import MyModel, NewManager
reveal_type(NewManager()) # N: Revealed type is "myapp.models.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) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
reveal_type(MyModel.objects.get()) # N: Revealed type is "Any" 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_only_method()) # N: Revealed type is "builtins.int"
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str" reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str"
@@ -74,7 +123,7 @@
main: | main: |
from myapp.models import MyModel, NewManager from myapp.models import MyModel, NewManager
reveal_type(NewManager()) # N: Revealed type is "myapp.models.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) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
reveal_type(MyModel.objects.get()) # N: Revealed type is "Any" 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_only_method()) # N: Revealed type is "builtins.int"
reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str" reveal_type(MyModel.objects.manager_and_queryset_method()) # N: Revealed type is "builtins.str"
@@ -187,6 +236,7 @@
from myapp.models import MyModel from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]" reveal_type(MyModel().objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str" reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is "builtins.str"
reveal_type(MyModel.objects.queryset_method_2()) # N: Revealed type is "builtins.int"
installed_apps: installed_apps:
- myapp - myapp
files: files:
@@ -200,6 +250,62 @@
def queryset_method(self) -> str: def queryset_method(self) -> str:
return 'hello' return 'hello'
@transaction.atomic
@transaction.atomic
def queryset_method_2(self) -> int:
return 2
NewManager = models.Manager.from_queryset(ModelQuerySet) NewManager = models.Manager.from_queryset(ModelQuerySet)
class MyModel(models.Model): class MyModel(models.Model):
objects = NewManager() objects = NewManager()
- case: from_queryset_model_gets_generated_manager_as_default_manager
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
reveal_type(MyModel.objects.queryset_method()) # N: Revealed type is "builtins.str"
reveal_type(MyModel._default_manager) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
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_can_resolve_explicit_any_methods
main: |
from myapp.models import MyModel
reveal_type(MyModel.objects) # N: Revealed type is "myapp.models.NewManager[myapp.models.MyModel]"
reveal_type(MyModel.objects.queryset_method(1)) # N: Revealed type is "Any"
reveal_type(MyModel.objects.queryset_method) # N: Revealed type is "def (qarg: Any) -> Any"
reveal_type(MyModel.objects.manager_method(2)) # N: Revealed type is "Any"
reveal_type(MyModel.objects.manager_method) # N: Revealed type is "def (marg: Any) -> Any"
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from typing import Any
class ModelQuerySet(models.QuerySet):
def queryset_method(self, qarg: Any) -> Any:
return 'hello'
class MyManager(models.Manager):
def manager_method(self, marg: Any) -> Any:
return 'hello'
NewManager = MyManager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()