def add_new_class_for_current_module(self, name: str, bases: List[Instance]) -> TypeInfo: current_module = self.api.modules[self.model_classdef.info.module_name] new_class_info = helpers.add_new_class_for_module(current_module, name=name, bases=bases) return new_class_info
def run_with_model_cls(self, model_cls: Type[Model]) -> None: for manager_name, manager in model_cls._meta.managers_map.items(): manager_fullname = helpers.get_class_fullname(manager.__class__) manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname) if manager_name not in self.model_classdef.info.names: 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 has_manager_any_base = any(self._is_manager_any(base) for base in manager_info.bases) if has_manager_any_base: custom_model_manager_name = manager.model.__name__ + '_' + manager.__class__.__name__ bases = [] for original_base in manager_info.bases: if self._is_manager_any(original_base): if original_base.type is None: raise helpers.IncompleteDefnException() original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) current_module = self.api.modules[self.model_classdef.info.module_name] custom_manager_info = helpers.add_new_class_for_module(current_module, custom_model_manager_name, bases=bases, fields=OrderedDict()) custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, 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(): manager_fullname = helpers.get_class_fullname(manager.__class__) manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname) if manager_name not in self.model_classdef.info.names: manager_type = Instance(manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, manager_type) else: # create new MODELNAME_MANAGERCLASSNAME class that represents manager parametrized with current model has_manager_any_base = any(self._is_manager_any(base) for base in manager_info.bases) if has_manager_any_base: custom_model_manager_name = manager.model.__name__ + '_' + manager.__class__.__name__ bases = [] for original_base in manager_info.bases: if self._is_manager_any(original_base): if original_base.type is None: if not self.api.final_iteration: self.api.defer() original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) current_module = self.api.modules[self.model_classdef.info.module_name] custom_manager_info = helpers.add_new_class_for_module(current_module, custom_model_manager_name, bases=bases, fields=OrderedDict()) custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) self.add_new_node_to_model_class(manager_name, custom_manager_type) # add _default_manager if '_default_manager' not in self.model_classdef.info.names: default_manager_fullname = helpers.get_class_fullname(model_cls._meta.default_manager.__class__) default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(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) # add related managers for relation in self.django_context.get_model_relations(model_cls): attname = relation.get_accessor_name() if attname is None: # no reverse accessor continue related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error(relation.related_model) if isinstance(relation, OneToOneRel): self.add_new_node_to_model_class(attname, Instance(related_model_info, [])) continue if isinstance(relation, (ManyToOneRel, ManyToManyRel)): manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS_FULLNAME) self.add_new_node_to_model_class(attname, Instance(manager_info, [Instance(related_model_info, [])])) continue
def get_or_create_annotated_type( api: Union[SemanticAnalyzer, CheckerPluginInterface], model_type: Instance, fields_dict: Optional[TypedDictType]) -> Instance: """ Get or create the type for a model for which you getting/setting any attr is allowed. The generated type is an subclass of the model and django._AnyAttrAllowed. The generated type is placed in the django_stubs_ext module, with the name WithAnnotations[ModelName]. If the user wanted to annotate their code using this type, then this is the annotation they would use. This is a bit of a hack to make a pretty type for error messages and which would make sense for users. """ model_module_name = "django_stubs_ext" if helpers.is_annotated_model_fullname(model_type.type.fullname): # If it's already a generated class, we want to use the original model as a base model_type = model_type.type.bases[0] if fields_dict is not None: type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}, {fields_dict}]" else: type_name = f"WithAnnotations[{model_type.type.fullname.replace('.', '__')}]" annotated_typeinfo = helpers.lookup_fully_qualified_typeinfo( cast(TypeChecker, api), model_module_name + "." + type_name) if annotated_typeinfo is None: model_module_file = api.modules[model_module_name] # type: ignore if isinstance(api, SemanticAnalyzer): annotated_model_type = api.named_type_or_none( ANY_ATTR_ALLOWED_CLASS_FULLNAME, []) assert annotated_model_type is not None else: annotated_model_type = api.named_generic_type( ANY_ATTR_ALLOWED_CLASS_FULLNAME, []) annotated_typeinfo = add_new_class_for_module( model_module_file, type_name, bases=[model_type] if fields_dict is not None else [model_type, annotated_model_type], fields=fields_dict.items if fields_dict is not None else None, no_serialize=True, ) if fields_dict is not None: # To allow structural subtyping, make it a Protocol annotated_typeinfo.is_protocol = True # Save for later to easily find which field types were annotated annotated_typeinfo.metadata[ "annotated_field_types"] = fields_dict.items annotated_type = Instance(annotated_typeinfo, []) return annotated_type
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): if original_base.type is None: raise helpers.IncompleteDefnException() original_base = helpers.reparametrize_instance(original_base, [Instance(self.model_classdef.info, [])]) bases.append(original_base) current_module = self.api.modules[self.model_classdef.info.module_name] custom_manager_info = helpers.add_new_class_for_module(current_module, name=name, bases=bases) # copy fields to a new manager new_cls_def_context = ClassDefContext(cls=custom_manager_info.defn, reason=self.ctx.reason, api=self.api) custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])]) for name, sym in base_manager_info.names.items(): # replace self type with new class, if copying method if isinstance(sym.node, FuncDef): helpers.copy_method_to_another_class(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 = custom_manager_info new_var._fullname = custom_manager_info.fullname + '.' + name new_sym.node = new_var custom_manager_info.names[name] = new_sym return custom_manager_type
def run_with_model_cls(self, model_cls: Type[Model]) -> None: # add related managers for relation in self.django_context.get_model_relations(model_cls): attname = relation.get_accessor_name() if attname is None or attname in self.model_classdef.info.names: # No reverse accessor or already declared. Note that this would also leave any # explicitly declared(i.e. non-inferred) reverse accessors alone continue related_model_cls = self.django_context.get_field_related_model_cls( relation) if related_model_cls is None: continue try: related_model_info = self.lookup_class_typeinfo_or_incomplete_defn_error( related_model_cls) except helpers.IncompleteDefnException as exc: if not self.api.final_iteration: raise exc else: continue if isinstance(relation, OneToOneRel): self.add_new_node_to_model_class( attname, Instance(related_model_info, [])) continue if isinstance(relation, (ManyToOneRel, ManyToManyRel)): try: related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error( fullnames.RELATED_MANAGER_CLASS) # noqa: E501 default_manager = related_model_info.names.get( "_default_manager") if not default_manager: raise helpers.IncompleteDefnException() except helpers.IncompleteDefnException as exc: if not self.api.final_iteration: raise exc else: continue # Check if the related model has a related manager subclassed from the default manager # TODO: Support other reverse managers than `_default_manager` default_reverse_manager_info = self.get_reverse_manager_info( model_info=related_model_info, derived_from="_default_manager") if default_reverse_manager_info: self.add_new_node_to_model_class( attname, Instance(default_reverse_manager_info, []), no_serialize=True) continue # The reverse manager we're looking for doesn't exist. So we create it. # The (default) reverse manager type is built from a RelatedManager and the default manager on the related model parametrized_related_manager_type = Instance( related_manager_info, [Instance(related_model_info, [])]) default_manager_type = default_manager.type assert default_manager_type is not None assert isinstance(default_manager_type, Instance) # When the default manager isn't custom there's no need to create a new type # as `RelatedManager` has `models.Manager` as base if default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME: self.add_new_node_to_model_class( attname, parametrized_related_manager_type) continue # The reverse manager is based on the related model's manager, so it makes most sense to add the new # related manager in that module new_related_manager_info = helpers.add_new_class_for_module( module=self.api.modules[related_model_info.module_name], name=f"{related_model_cls.__name__}_RelatedManager", bases=[ parametrized_related_manager_type, default_manager_type ], no_serialize=True, ) new_related_manager_info.metadata["django"] = { "related_manager_to_model": related_model_info.fullname } # Stash the new reverse manager type fullname on the related model, so we don't duplicate # or have to create it again for other reverse relations self.set_reverse_manager_info( related_model_info, derived_from="_default_manager", fullname=new_related_manager_info.fullname, ) self.add_new_node_to_model_class(attname, Instance( new_related_manager_info, []), no_serialize=True)