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