from typing import Final from mypy.checker import TypeChecker from mypy.copytype import copy_type from mypy.nodes import ( GDEF, CallExpr, Decorator, FuncBase, FuncDef, MemberExpr, Node, OverloadedFuncDef, PlaceholderNode, RefExpr, StrExpr, SymbolTableNode, TypeInfo, ) from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext from mypy.plugins.common import add_method_to_class from mypy.semanal import SemanticAnalyzer from mypy.semanal_shared import has_placeholder from mypy.subtypes import find_member from mypy.types import ( AnyType, CallableType, FunctionLike, Instance, Overloaded, ProperType, TupleType, TypeOfAny, TypeType, TypeVarType, UnionType, get_proper_type, ) from mypy.types import Type as MypyType from mypy.typevars import fill_typevars from mypy_django_plugin.lib import fullnames, helpers MANAGER_METHODS_RETURNING_QUERYSET: Final = frozenset( ( "alias", "all", "annotate", "complex_filter", "defer", "difference", "distinct", "exclude", "extra", "filter", "intersection", "none", "only", "order_by", "prefetch_related", "reverse", "select_for_update", "select_related", "union", "using", ) ) def get_method_type_from_dynamic_manager( api: TypeChecker, method_name: str, manager_instance: Instance ) -> ProperType | None: """ Attempt to resolve a method on a manager that was built from '.from_queryset' """ manager_type_info = manager_instance.type.get_containing_type_info(method_name) if ( manager_type_info is None or "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 is_fallback_queryset = queryset_info.metadata.get("django", {}).get("any_fallback_queryset", False) base_that_has_method = queryset_info.get_containing_type_info(method_name) if base_that_has_method is None: return None method_type = _get_funcdef_type(base_that_has_method.names[method_name].node) if not isinstance(method_type, FunctionLike): return method_type items = [] for item in method_type.items: items.append( _process_dynamic_method( method_name, item, base_that_has_method=base_that_has_method, queryset_info=queryset_info, manager_instance=manager_instance, is_fallback_queryset=is_fallback_queryset, ) ) return Overloaded(items) if len(items) > 1 else items[0] def _process_dynamic_method( method_name: str, method_type: CallableType, *, queryset_info: TypeInfo, base_that_has_method: TypeInfo, manager_instance: Instance, is_fallback_queryset: bool, ) -> CallableType: variables = method_type.variables ret_type = method_type.ret_type manager_model = get_proper_type(find_member("model", manager_instance, manager_instance)) assert isinstance(manager_model, TypeType), manager_model manager_model_type = manager_model.item queryset_instance = Instance(queryset_info, (manager_model_type,) * len(queryset_info.type_vars)) # For methods on the manager that return a queryset we need to override the # return type to be the actual queryset class, not the base QuerySet that's # used by the typing stubs. if method_name in MANAGER_METHODS_RETURNING_QUERYSET: ret_type = queryset_instance variables = () args_types = method_type.arg_types[1:] if _has_compatible_type_vars(base_that_has_method): typed_var = manager_instance.args or queryset_info.bases[0].args if ( typed_var and isinstance((model_arg := get_proper_type(typed_var[0])), Instance) and model_arg.type.has_base(fullnames.MODEL_CLASS_FULLNAME) and helpers.is_model_type(model_arg.type) ): ret_type = _replace_type_var(ret_type, base_that_has_method.defn.type_vars[0].fullname, model_arg) args_types = [ _replace_type_var(arg_type, base_that_has_method.defn.type_vars[0].fullname, model_arg) for arg_type in args_types ] if base_that_has_method.self_type: # Manages -> Self returns ret_type = _replace_type_var(ret_type, base_that_has_method.self_type.fullname, queryset_instance) # Drop any 'self' argument as our manager is already initialized return method_type.copy_modified( arg_types=args_types, arg_kinds=method_type.arg_kinds[1:], arg_names=method_type.arg_names[1:], variables=variables, ret_type=ret_type, ) def _get_funcdef_type(definition: Node | None) -> ProperType | None: if isinstance(definition, FuncBase): return definition.type if isinstance(definition, Decorator): return definition.func.type return None def _has_compatible_type_vars(type_info: TypeInfo) -> bool: """ Determines whether the provided 'type_info', is a generically parameterized subclass of models.QuerySet[T], with exactly one type variable. Criteria for compatibility: 1. 'type_info' must be a generic class with exactly one type variable. 2. All superclasses of 'type_info', up to and including models.QuerySet, must also be generic classes with exactly one type variable. Examples: Compatible: class _MyModelQuerySet(models.QuerySet[T]): ... class MyModelQuerySet(_MyModelQuerySet[T_2]): def example(self) -> T_2: ... Incompatible: class MyModelQuerySet(models.QuerySet[T], Generic[T, T2]): def example(self, a: T2) -> T_2: ... Returns: True if the 'base' meets the criteria, otherwise False. """ args = type_info.defn.type_vars if not args or len(args) > 1 or not isinstance(args[0], TypeVarType): # No type var to manage, or too many return False type_var: MypyType | None = None # check that for all the bases it has only one type vars for sub_base in type_info.bases: unic_args = list(set(sub_base.args)) if not unic_args or len(unic_args) > 1: # No type var for the sub_base, skipping continue if type_var and unic_args and type_var != unic_args[0]: # There is two different type vars in the bases, we are not compatible return False type_var = unic_args[0] if not type_var: # No type var found in the bases. return False if type_info.has_base(fullnames.QUERYSET_CLASS_FULLNAME): # If it is a subclass of QuerySet, it is compatible. return True # check that at least one base is a subclass of queryset with Generic type vars return any(_has_compatible_type_vars(sub_base.type) for sub_base in type_info.bases) def _replace_type_var(ret_type: MypyType, to_replace: str, replace_by: MypyType) -> MypyType: """ Substitutes a specified type variable within a Mypy type expression with an actual type. This function is recursive, and it operates on various kinds of Mypy types like Instance, ProperType, etc., to deeply replace the specified type variable. Parameters: - ret_type: A Mypy type expression where the substitution should occur. - to_replace: The type variable to be replaced, specified as its full name. - replace_by: The actual Mypy type to substitute in place of 'to_replace'. Example: Given: ret_type = "typing.Collection[T]" to_replace = "T" replace_by = "myapp.models.MyModel" Result: "typing.Collection[myapp.models.MyModel]" """ if isinstance(ret_type, TypeVarType) and ret_type.fullname == to_replace: return replace_by ret_type = copy_type(get_proper_type(ret_type)) if isinstance(ret_type, Instance): # Since it is an instance, recursively find the type var for all its args. return ret_type.copy_modified( args=tuple(_replace_type_var(item, to_replace, replace_by) for item in ret_type.args) ) if isinstance(ret_type, TypeType): return TypeType.make_normalized( _replace_type_var(ret_type.item, to_replace, replace_by), line=ret_type.line, column=ret_type.column, ) if isinstance(ret_type, TupleType): return ret_type.copy_modified( items=[_replace_type_var(item, to_replace, replace_by) for item in ret_type.items] ) if isinstance(ret_type, UnionType): return UnionType.make_union( items=[_replace_type_var(item, to_replace, replace_by) for item in ret_type.items], line=ret_type.line, column=ret_type.column, ) return ret_type def resolve_manager_method_from_instance(instance: Instance, method_name: str, ctx: AttributeContext) -> MypyType: api = helpers.get_typechecker_api(ctx) method_type = get_method_type_from_dynamic_manager(api, method_name, instance) return method_type if method_type is not None else ctx.default_attr_type def resolve_manager_method(ctx: AttributeContext) -> 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. """ # Skip (method) type that is currently something other than Any of type `implementation_artifact` default_attr_type = get_proper_type(ctx.default_attr_type) if not isinstance(default_attr_type, AnyType): return ctx.default_attr_type if default_attr_type.type_of_any != TypeOfAny.implementation_artifact: 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. if isinstance(ctx.context, MemberExpr): method_name = ctx.context.name elif isinstance(ctx.context, CallExpr) and isinstance(ctx.context.callee, MemberExpr): method_name = ctx.context.callee.name else: ctx.api.fail("Unable to resolve return type of queryset/manager method", ctx.context) return AnyType(TypeOfAny.from_error) if isinstance(ctx.type, Instance): return resolve_manager_method_from_instance(instance=ctx.type, method_name=method_name, ctx=ctx) if isinstance(ctx.type, UnionType) and all(isinstance(get_proper_type(item), Instance) for item in ctx.type.items): resolved = tuple( resolve_manager_method_from_instance(instance=instance, method_name=method_name, ctx=ctx) for item in ctx.type.items if isinstance((instance := get_proper_type(item)), Instance) ) return UnionType(resolved) ctx.api.fail(f'Unable to resolve return type of queryset/manager method "{method_name}"', ctx.context) return AnyType(TypeOfAny.from_error) def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None: """ Insert a new manager class node for a: ' = .from_queryset()'. When the assignment expression lives at module level. """ semanal_api = helpers.get_semanal_api(ctx) # TODO: Emit an error when called in a class scope if semanal_api.is_class_scope(): return # Don't redeclare the manager class if we've already defined it. manager_sym = semanal_api.lookup_current_scope(ctx.name) if manager_sym and isinstance(manager_sym.node, TypeInfo): # This is just a deferral run where our work is already finished return new_manager_info = create_manager_info_from_from_queryset_call(semanal_api, ctx.call, ctx.name) if new_manager_info is None: if not ctx.api.final_iteration: # Only add PlaceholderNode if it doesn't already exist to prevent # infinite semantic analysis iterations (fixes #2373) if not (manager_sym and isinstance(manager_sym.node, PlaceholderNode)): # XXX: hack for python/mypy#17402 ph = PlaceholderNode(ctx.api.qualified_name(ctx.name), ctx.call, ctx.call.line, becomes_typeinfo=True) ctx.api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, ph)) ctx.api.defer() return def register_dynamically_created_manager(fullname: str, manager_name: str, manager_base: TypeInfo) -> None: manager_base.metadata.setdefault("from_queryset_managers", {}) # The `__module__` value of the manager type created by Django's # `.from_queryset` is `django.db.models.manager`. But we put new type(s) in the # module currently being processed, so we'll map those together through metadata. runtime_fullname = ".".join(["django.db.models.manager", manager_name]) manager_base.metadata["from_queryset_managers"][runtime_fullname] = fullname def create_manager_info_from_from_queryset_call( api: SemanticAnalyzer, call_expr: CallExpr, name: str | None = None ) -> TypeInfo | None: """ Extract manager and queryset TypeInfo from a from_queryset call. """ if ( # Check that this is a from_queryset call on a manager subclass not isinstance(call_expr.callee, MemberExpr) or not isinstance(call_expr.callee.expr, RefExpr) or not isinstance(call_expr.callee.expr.node, TypeInfo) or not call_expr.callee.expr.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME) or not call_expr.callee.name == "from_queryset" # Check that the call has one or two arguments and that the first is a # QuerySet subclass or not 1 <= len(call_expr.args) <= 2 or not isinstance(call_expr.args[0], RefExpr) or not isinstance(call_expr.args[0].node, TypeInfo) or not call_expr.args[0].node.has_base(fullnames.QUERYSET_CLASS_FULLNAME) ): return None base_manager_info, queryset_info = call_expr.callee.expr.node, call_expr.args[0].node if queryset_info.fullname is None: # In some cases, due to the way the semantic analyzer works, only # passed_queryset.name is available. But it should be analyzed again, # so this isn't a problem. return None # type: ignore[unreachable] if len(call_expr.args) == 2 and isinstance(call_expr.args[1], StrExpr): manager_name = call_expr.args[1].value else: manager_name = f"{base_manager_info.name}From{queryset_info.name}" # Always look in global scope, as that's where we'll declare dynamic manager classes manager_sym = api.globals.get(manager_name) if ( manager_sym is not None and isinstance(manager_sym.node, TypeInfo) and manager_sym.node.has_base(base_manager_info.fullname) and manager_sym.node.metadata.get("django", {}).get("from_queryset_manager") == queryset_info.fullname ): # Reuse an identical, already generated, manager new_manager_info = manager_sym.node else: # Create a new `TypeInfo` instance for the manager type try: new_manager_info = create_manager_class( api=api, base_manager_info=base_manager_info, name=manager_name, line=call_expr.line, with_unique_name=name is not None and name != manager_name, ) except helpers.IncompleteDefnException: return None populate_manager_from_queryset(new_manager_info, queryset_info) register_dynamically_created_manager( fullname=new_manager_info.fullname, manager_name=manager_name, manager_base=base_manager_info, ) # Add the new manager to the current module # TODO: use proper SemanticAnalyzer API for that. module = api.modules[api.cur_mod_id] if name is not None and name != new_manager_info.name: # Unless names are equal, there's 2 symbol names that needs the manager info module.names[name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) module.names[new_manager_info.name] = SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) return new_manager_info def create_manager_class( api: SemanticAnalyzer, base_manager_info: TypeInfo, name: str, line: int, with_unique_name: bool ) -> TypeInfo: base_manager_instance = fill_typevars(base_manager_info) assert isinstance(base_manager_instance, Instance) # If any of the type vars are undefined we need to defer. This is handled by the caller if any(has_placeholder(type_var) for type_var in base_manager_info.defn.type_vars): raise helpers.IncompleteDefnException if with_unique_name: manager_info = helpers.add_new_class_for_module( module=api.modules[api.cur_mod_id], name=name, bases=[base_manager_instance], ) else: manager_info = helpers.create_type_info(name, api.cur_mod_id, bases=[base_manager_instance]) manager_info.line = line manager_info.type_vars = base_manager_info.type_vars manager_info.defn.type_vars = base_manager_info.defn.type_vars manager_info.defn.line = line return manager_info def populate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeInfo) -> None: """ Add methods from the QuerySet class to the manager. """ # Stash the queryset fullname which was passed to .from_queryset So that # our 'resolve_manager_method' attribute hook can fetch the method from # that QuerySet class django_metadata = helpers.get_django_metadata(manager_info) django_metadata["from_queryset_manager"] = queryset_info.fullname # We collect and mark up all methods before django.db.models.query.QuerySet as class members for class_mro_info in queryset_info.mro: if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME: break for name, sym in class_mro_info.names.items(): if not isinstance(sym.node, FuncDef | OverloadedFuncDef | Decorator): continue # private, magic methods are not copied # https://github.com/django/django/blob/5.0.4/django/db/models/manager.py#L101 if name.startswith("_"): continue # 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( manager_info, name=name, sym_type=AnyType(TypeOfAny.implementation_artifact), ) # For methods on BaseManager that return a queryset we need to update # the return type to be the actual queryset subclass used. This is done # by adding the methods as attributes with type Any to the manager # class. The actual type of these methods are resolved in # resolve_manager_method. for method_name in MANAGER_METHODS_RETURNING_QUERYSET: helpers.add_new_sym_for_info( manager_info, name=method_name, sym_type=AnyType(TypeOfAny.implementation_artifact), ) def add_as_manager_to_queryset_class(ctx: ClassDefContext) -> None: semanal_api = helpers.get_semanal_api(ctx) def _defer() -> None: if not semanal_api.final_iteration: semanal_api.defer() queryset_info = semanal_api.type if queryset_info is None: return _defer() # either a manual `as_manager` definition or this is a deferral pass if "as_manager" in queryset_info.names: return base_as_manager = queryset_info.get("as_manager") if base_as_manager is None: return manager_sym = semanal_api.lookup_fully_qualified_or_none(fullnames.MANAGER_CLASS_FULLNAME) if manager_sym is None or not isinstance(manager_sym.node, TypeInfo): return _defer() manager_base = manager_sym.node base_ret_type = manager_base if base_as_manager.fullname != f"{fullnames.QUERYSET_CLASS_FULLNAME}.as_manager": base_as_manager_type = get_proper_type(base_as_manager.type) if not isinstance(base_as_manager_type, CallableType): return base_as_manager_ret_type = get_proper_type(base_as_manager_type.ret_type) if not isinstance(base_as_manager_ret_type, Instance): return base_ret_type = base_as_manager_ret_type.type manager_class_name = f"{manager_base.name}From{queryset_info.name}" current_module = semanal_api.modules[semanal_api.cur_mod_id] existing_sym = current_module.names.get(manager_class_name) if ( existing_sym is not None and isinstance(existing_sym.node, TypeInfo) and existing_sym.node.has_base(fullnames.MANAGER_CLASS_FULLNAME) and existing_sym.node.metadata.get("django", {}).get("from_queryset_manager") == queryset_info.fullname ): # Reuse an identical, already generated, manager new_manager_info = existing_sym.node else: # Create a new `TypeInfo` instance for the manager type try: new_manager_info = create_manager_class( api=semanal_api, base_manager_info=base_ret_type, name=manager_class_name, line=queryset_info.line, with_unique_name=True, ) except helpers.IncompleteDefnException: return _defer() populate_manager_from_queryset(new_manager_info, queryset_info) register_dynamically_created_manager( fullname=new_manager_info.fullname, manager_name=manager_class_name, manager_base=manager_base, ) # Add the new manager to the current module # We'll use `new_manager_info.name` instead of `manager_class_name` here # to handle possible name collisions, as it's unique. current_module.names[new_manager_info.name] = ( # Note that the generated manager type is always inserted at module level SymbolTableNode(GDEF, new_manager_info, plugin_generated=True) ) add_method_to_class( semanal_api, ctx.cls, "as_manager", args=[], return_type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)]), is_classmethod=True, ) def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None: """ Add implicit generics to manager classes that are defined without generic. Eg. class MyManager(models.Manager): ... is interpreted as: _T = TypeVar('_T', covariant=True) class MyManager(models.Manager[_T]): ... Note that this does not happen if mypy is run with disallow_any_generics = True, as not specifying the generic type is then considered an error. """ manager = ctx.api.lookup_fully_qualified_or_none(ctx.cls.fullname) if manager is None or manager.node is None: return assert isinstance(manager.node, TypeInfo) if manager.node.type_vars: # We've already been here return parent_manager = next( (base for base in manager.node.bases if base.type.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)), None, ) if parent_manager is None or len(parent_manager.args) != 1: return model_param = get_proper_type(parent_manager.args[0]) if not isinstance(model_param, AnyType) or model_param.type_of_any is not TypeOfAny.from_omitted_generics: return type_vars = tuple(parent_manager.type.defn.type_vars) # If we end up with placeholders we need to defer so the placeholders are # resolved in a future iteration if any(has_placeholder(type_var) for type_var in type_vars): if not ctx.api.final_iteration: ctx.api.defer() else: return parent_manager.args = type_vars manager.node.defn.type_vars = list(type_vars) manager.node.add_type_vars()