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)
