def add_new_node_to_model_class(self,
name: str,
typ: MypyType,
no_serialize: bool = False) -> None:
helpers.add_new_sym_for_info(self.model_classdef.info,
name=name,
sym_type=typ,
no_serialize=no_serialize)
def fill_descriptor_types_for_related_field(ctx: FunctionContext, django_context: DjangoContext) -> MypyType:
current_field = _get_current_field_from_assignment(ctx, django_context)
if current_field is None:
return AnyType(TypeOfAny.from_error)
assert isinstance(current_field, RelatedField)
related_model_cls = django_context.get_field_related_model_cls(current_field)
if related_model_cls is None:
return AnyType(TypeOfAny.from_error)
default_related_field_type = set_descriptor_types_for_field(ctx)
# self reference with abstract=True on the model where ForeignKey is defined
current_model_cls = current_field.model
if (current_model_cls._meta.abstract
and current_model_cls == related_model_cls):
# for all derived non-abstract classes, set variable with this name to
# __get__/__set__ of ForeignKey of derived model
for model_cls in django_context.all_registered_model_classes:
if issubclass(model_cls, current_model_cls) and not model_cls._meta.abstract:
derived_model_info = helpers.lookup_class_typeinfo(helpers.get_typechecker_api(ctx), model_cls)
if derived_model_info is not None:
fk_ref_type = Instance(derived_model_info, [])
derived_fk_type = reparametrize_related_field_type(default_related_field_type,
set_type=fk_ref_type, get_type=fk_ref_type)
helpers.add_new_sym_for_info(derived_model_info,
name=current_field.name,
sym_type=derived_fk_type)
related_model = related_model_cls
related_model_to_set = related_model_cls
if related_model_to_set._meta.proxy_for_model is not None:
related_model_to_set = related_model_to_set._meta.proxy_for_model
typechecker_api = helpers.get_typechecker_api(ctx)
related_model_info = helpers.lookup_class_typeinfo(typechecker_api, related_model)
if related_model_info is None:
# maybe no type stub
related_model_type = AnyType(TypeOfAny.unannotated)
else:
related_model_type = Instance(related_model_info, []) # type: ignore
related_model_to_set_info = helpers.lookup_class_typeinfo(typechecker_api, related_model_to_set)
if related_model_to_set_info is None:
# maybe no type stub
related_model_to_set_type = AnyType(TypeOfAny.unannotated)
else:
related_model_to_set_type = Instance(related_model_to_set_info, []) # type: ignore
# replace Any with referred_to_type
return reparametrize_related_field_type(default_related_field_type,
set_type=related_model_to_set_type,
get_type=related_model_type)
def add_new_node_to_model_class(self, name: str, typ: MypyType) -> None:
helpers.add_new_sym_for_info(self.model_classdef.info,
name=name,
sym_type=typ)
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)
# 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
assert isinstance(callee, MemberExpr)
assert isinstance(callee.expr, RefExpr)
base_manager_info = callee.expr.node
if base_manager_info is None:
if not semanal_api.final_iteration:
semanal_api.defer()
return
assert isinstance(base_manager_info, TypeInfo)
passed_queryset = ctx.call.args[0]
assert isinstance(passed_queryset, NameExpr)
derived_queryset_fullname = passed_queryset.fullname
if derived_queryset_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
base_manager_instance = fill_typevars(base_manager_info)
assert isinstance(base_manager_instance, Instance)
new_manager_info = semanal_api.basic_new_typeinfo(
ctx.name,
basetype_or_fallback=base_manager_instance,
line=ctx.call.line)
sym = semanal_api.lookup_fully_qualified_or_none(derived_queryset_fullname)
assert sym is not None
if sym.node is None:
if not semanal_api.final_iteration:
semanal_api.defer()
else:
# inherit from Any to prevent false-positives, if queryset class cannot be resolved
new_manager_info.fallback_to_any = True
return
derived_queryset_info = sym.node
assert isinstance(derived_queryset_info, TypeInfo)
new_manager_info.line = ctx.call.line
new_manager_info.type_vars = base_manager_info.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.metaclass_type = new_manager_info.calculate_metaclass_type(
)
# 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
new_manager_info.metadata["django"] = {
"from_queryset_manager": derived_queryset_fullname
}
if len(ctx.call.args) > 1:
expr = ctx.call.args[1]
assert isinstance(expr, StrExpr)
custom_manager_generated_name = expr.value
else:
custom_manager_generated_name = base_manager_info.name + "From" + derived_queryset_info.name
custom_manager_generated_fullname = ".".join(
["django.db.models.manager", custom_manager_generated_name])
base_manager_info.metadata.setdefault("from_queryset_managers", {})
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
helpers.add_new_manager_base(semanal_api, new_manager_info.fullname)
class_def_context = ClassDefContext(cls=new_manager_info.defn,
reason=ctx.call,
api=semanal_api)
self_type = fill_typevars(new_manager_info)
assert isinstance(self_type, Instance)
# 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:
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, Decorator)):
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(
new_manager_info,
name=name,
sym_type=AnyType(TypeOfAny.special_form),
no_serialize=True,
)
# we need to copy all methods in MRO before django.db.models.query.QuerySet
# Gather names of all BaseManager methods
manager_method_names = []
for manager_mro_info in new_manager_info.mro:
if manager_mro_info.fullname == fullnames.BASE_MANAGER_CLASS_FULLNAME:
for name, sym in manager_mro_info.names.items():
manager_method_names.append(name)
# Copy/alter all methods in common between BaseManager/QuerySet over to the new manager if their return type is
# the QuerySet's self-type. Alter the return type to be the custom queryset, parameterized by the manager's model
# type variable.
for class_mro_info in derived_queryset_info.mro:
if class_mro_info.fullname != fullnames.QUERYSET_CLASS_FULLNAME:
continue
for name, sym in class_mro_info.names.items():
if name not in manager_method_names:
continue
if isinstance(sym.node, FuncDef):
func_node = sym.node
elif isinstance(sym.node, Decorator):
func_node = sym.node.func
else:
continue
method_type = func_node.type
if not isinstance(method_type, CallableType):
if not semanal_api.final_iteration:
semanal_api.defer()
return None
original_return_type = method_type.ret_type
if original_return_type is None:
continue
# Skip any method that doesn't return _QS
original_return_type = get_proper_type(original_return_type)
if isinstance(original_return_type, UnboundType):
if original_return_type.name != "_QS":
continue
elif isinstance(original_return_type, TypeVarType):
if original_return_type.name != "_QS":
continue
else:
continue
# Return the custom queryset parameterized by the manager's type vars
return_type = Instance(derived_queryset_info, self_type.args)
helpers.copy_method_to_another_class(
class_def_context,
self_type,
new_method_name=name,
method_node=func_node,
return_type=return_type,
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,
no_serialize=True))
