mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-08 04:54:48 +08:00
70 lines
3.1 KiB
Python
70 lines
3.1 KiB
Python
from mypy.checker import gen_unique_name
|
|
from mypy.plugin import AttributeContext
|
|
from mypy.types import Instance
|
|
from mypy.types import Type as MypyType
|
|
|
|
from django.db.models.fields.reverse_related import ForeignObjectRel, OneToOneRel, ManyToOneRel, ManyToManyRel
|
|
|
|
from mypy_django_plugin.lib import helpers, fullnames
|
|
from mypy_django_plugin.lib.helpers import GetAttributeCallback
|
|
|
|
|
|
class GetRelatedManagerCallback(GetAttributeCallback):
|
|
obj_type: Instance
|
|
|
|
def get_related_manager_type(self, relation: ForeignObjectRel) -> MypyType:
|
|
related_model_cls = self.django_context.get_field_related_model_cls(relation)
|
|
if related_model_cls is None:
|
|
# could not find a referenced model (maybe invalid to= value, or GenericForeignKey)
|
|
# TODO: show error
|
|
return self.default_attr_type
|
|
|
|
related_model_info = self.lookup_typeinfo(helpers.get_class_fullname(related_model_cls))
|
|
if related_model_info is None:
|
|
# TODO: show error
|
|
return self.default_attr_type
|
|
|
|
if isinstance(relation, OneToOneRel):
|
|
return Instance(related_model_info, [])
|
|
|
|
elif isinstance(relation, (ManyToOneRel, ManyToManyRel)):
|
|
related_manager_info = self.lookup_typeinfo(fullnames.RELATED_MANAGER_CLASS)
|
|
if related_manager_info is None:
|
|
return self.default_attr_type
|
|
|
|
# get type of default_manager for model
|
|
default_manager_fullname = helpers.get_class_fullname(related_model_cls._meta.default_manager.__class__)
|
|
default_manager_info = self.lookup_typeinfo(default_manager_fullname)
|
|
if default_manager_info is None:
|
|
return self.default_attr_type
|
|
|
|
default_manager_type = Instance(default_manager_info, [Instance(related_model_info, [])])
|
|
related_manager_type = Instance(related_manager_info,
|
|
[Instance(related_model_info, [])])
|
|
|
|
if (not isinstance(default_manager_type, Instance)
|
|
or default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME):
|
|
# if not defined or trivial -> just return RelatedManager[Model]
|
|
return related_manager_type
|
|
|
|
# make anonymous class
|
|
name = gen_unique_name(related_model_cls.__name__ + '_' + 'RelatedManager',
|
|
self.obj_type.type.names)
|
|
bases = [related_manager_type, default_manager_type]
|
|
new_manager_info = self.new_typeinfo(name, bases)
|
|
return Instance(new_manager_info, [])
|
|
|
|
def __call__(self, ctx: AttributeContext):
|
|
super().__call__(ctx)
|
|
assert isinstance(self.obj_type, Instance)
|
|
|
|
model_fullname = self.obj_type.type.fullname
|
|
model_cls = self.django_context.get_model_class_by_fullname(model_fullname)
|
|
if model_cls is None:
|
|
return self.default_attr_type
|
|
for reverse_manager_name, relation in self.django_context.get_model_relations(model_cls):
|
|
if reverse_manager_name == self.name:
|
|
return self.get_related_manager_type(relation)
|
|
|
|
return self.default_attr_type
|