django-stubs/mypy_django_plugin_newsemanal/transformers/models.py

import dataclasses
from abc import ABCMeta, abstractmethod
from typing import cast

from django.db.models.fields.related import ForeignKey
from mypy.newsemanal.semanal import NewSemanticAnalyzer
from mypy.nodes import ClassDef, MDEF, SymbolTableNode, TypeInfo, Var
from mypy.plugin import ClassDefContext
from mypy.types import Instance

from django.db.models.fields import Field
from django.db.models.fields.reverse_related import ManyToOneRel, OneToOneRel, ManyToManyRel
from mypy_django_plugin_newsemanal.django.context import DjangoContext
from mypy_django_plugin_newsemanal.lib import helpers, fullnames
from mypy_django_plugin_newsemanal.transformers import fields
from mypy_django_plugin_newsemanal.transformers.fields import get_field_descriptor_types


@dataclasses.dataclass
class ModelClassInitializer(metaclass=ABCMeta):
    api: NewSemanticAnalyzer
    model_classdef: ClassDef
    django_context: DjangoContext
    ctx: ClassDefContext

    @classmethod
    def from_ctx(cls, ctx: ClassDefContext, django_context: DjangoContext):
        return cls(api=cast(NewSemanticAnalyzer, ctx.api),
                   model_classdef=ctx.cls,
                   django_context=django_context,
                   ctx=ctx)

    def lookup_typeinfo_or_incomplete_defn_error(self, fullname: str) -> TypeInfo:
        sym = self.api.lookup_fully_qualified_or_none(fullname)
        if sym is None or not isinstance(sym.node, TypeInfo):
            raise helpers.IncompleteDefnException(f'No {fullname!r} found')
        return sym.node

    def lookup_class_typeinfo_or_incomplete_defn_error(self, klass: type) -> TypeInfo:
        fullname = helpers.get_class_fullname(klass)
        field_info = self.lookup_typeinfo_or_incomplete_defn_error(fullname)
        return field_info

    def add_new_node_to_model_class(self, name: str, typ: Instance) -> None:
        # type=: type of the variable itself
        var = Var(name=name, type=typ)
        # var.info: type of the object variable is bound to
        var.info = self.model_classdef.info
        var._fullname = self.model_classdef.info.fullname() + '.' + name
        var.is_initialized_in_class = True
        var.is_inferred = True
        self.model_classdef.info.names[name] = SymbolTableNode(MDEF, var, plugin_generated=True)

    @abstractmethod
    def run(self) -> None:
        raise NotImplementedError()


class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
    """
    Replaces
        class MyModel(models.Model):
            class Meta:
                pass
    with
        class MyModel(models.Model):
            class Meta(Any):
                pass
    to get around incompatible Meta inner classes for different models.
    """

    def run(self) -> None:
        meta_node = helpers.get_nested_meta_node_for_current_class(self.model_classdef.info)
        if meta_node is None:
            return None
        meta_node.fallback_to_any = True


class AddDefaultPrimaryKey(ModelClassInitializer):
    def run(self) -> None:
        model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
        if model_cls is None:
            return

        auto_field = model_cls._meta.auto_field
        if auto_field and not self.model_classdef.info.has_readable_member(auto_field.attname):
            # autogenerated field
            auto_field_fullname = helpers.get_class_fullname(auto_field.__class__)
            auto_field_info = self.lookup_typeinfo_or_incomplete_defn_error(auto_field_fullname)

            set_type, get_type = fields.get_field_descriptor_types(auto_field_info, is_nullable=False)
            self.add_new_node_to_model_class(auto_field.attname, Instance(auto_field_info,
                                                                          [set_type, get_type]))


class AddRelatedModelsId(ModelClassInitializer):
    def run(self) -> None:
        model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
        if model_cls is None:
            return

        for field in model_cls._meta.get_fields():
            if isinstance(field, ForeignKey):
                rel_primary_key_field = self.django_context.get_primary_key_field(field.related_model)
                field_info = self.lookup_class_typeinfo_or_incomplete_defn_error(rel_primary_key_field.__class__)
                is_nullable = self.django_context.fields_context.get_field_nullability(field, None)
                set_type, get_type = get_field_descriptor_types(field_info, is_nullable)
                self.add_new_node_to_model_class(field.attname,
                                                 Instance(field_info, [set_type, get_type]))


class AddManagers(ModelClassInitializer):
    def run(self):
        model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
        if model_cls is None:
            return

        for manager_name, manager in model_cls._meta.managers_map.items():
            if manager_name not in self.model_classdef.info.names:
                manager_fullname = helpers.get_class_fullname(manager.__class__)
                manager_info = self.lookup_typeinfo_or_incomplete_defn_error(manager_fullname)

                manager = Instance(manager_info, [Instance(self.model_classdef.info, [])])
                self.add_new_node_to_model_class(manager_name, manager)

        # 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.related_name
            if attname is None:
                attname = relation.name + '_set'

            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 process_model_class(ctx: ClassDefContext,
                        django_context: DjangoContext) -> None:
    initializers = [
        InjectAnyAsBaseForNestedMeta,
        AddDefaultPrimaryKey,
        AddRelatedModelsId,
        AddManagers,
    ]
    for initializer_cls in initializers:
        try:
            initializer_cls.from_ctx(ctx, django_context).run()
        except helpers.IncompleteDefnException:
            if not ctx.api.final_iteration:
                ctx.api.defer()