wip

mypy_django_plugin/transformers2/related_managers.py

@@ -0,0 +1,69 @@
+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