remove -semanal suffix

mypy_django_plugin/transformers/models.py

@@ -0,0 +1,162 @@
+import dataclasses
+from abc import ABCMeta, abstractmethod
+from typing import cast
+
+from django.db.models.fields.related import ForeignKey
+from django.db.models.fields.reverse_related import ManyToManyRel, ManyToOneRel, OneToOneRel
+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 mypy_django_plugin.django.context import DjangoContext
+from mypy_django_plugin.lib import fullnames, helpers
+from mypy_django_plugin.transformers import fields
+from mypy_django_plugin.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()