new semanal wip 1

mypy_django_plugin/transformers/fields.py

@@ -1,13 +1,13 @@
 from typing import Optional, cast
 
 from mypy.checker import TypeChecker
-from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo
+from mypy.nodes import ListExpr, NameExpr, StrExpr, TupleExpr, TypeInfo, Expression
 from mypy.plugin import FunctionContext
 from mypy.types import (
     AnyType, CallableType, Instance, TupleType, Type, UnionType,
 )
 
-from mypy_django_plugin.lib import metadata, fullnames, helpers
+from mypy_django_plugin.lib import fullnames, helpers, metadata
 
 
 def extract_referred_to_type(ctx: FunctionContext) -> Optional[Instance]:
@@ -90,7 +90,7 @@ def fill_descriptor_types_for_related_field(ctx: FunctionContext) -> Type:
 
 def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance:
     default_return_type = cast(Instance, ctx.default_return_type)
-    is_nullable = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'null'))
+    is_nullable = helpers.parse_bool(helpers.get_call_argument_by_name(ctx, 'null'))
     set_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_set_type',
                                                    is_nullable=is_nullable)
     get_type = helpers.get_private_descriptor_type(default_return_type.type, '_pyi_private_get_type',
@@ -101,7 +101,7 @@ def set_descriptor_types_for_field(ctx: FunctionContext) -> Instance:
 def determine_type_of_array_field(ctx: FunctionContext) -> Type:
     default_return_type = set_descriptor_types_for_field(ctx)
 
-    base_field_arg_type = helpers.get_argument_type_by_name(ctx, 'base_field')
+    base_field_arg_type = helpers.get_call_argument_type_by_name(ctx, 'base_field')
     if not base_field_arg_type or not isinstance(base_field_arg_type, Instance):
         return default_return_type
 
@@ -118,9 +118,7 @@ def transform_into_proper_return_type(ctx: FunctionContext) -> Type:
     if not isinstance(default_return_type, Instance):
         return default_return_type
 
-    if helpers.has_any_of_bases(default_return_type.type, (fullnames.FOREIGN_KEY_FULLNAME,
-                                                           fullnames.ONETOONE_FIELD_FULLNAME,
-                                                           fullnames.MANYTOMANY_FIELD_FULLNAME)):
+    if helpers.has_any_of_bases(default_return_type.type, fullnames.RELATED_FIELDS_CLASSES):
         return fill_descriptor_types_for_related_field(ctx)
 
     if default_return_type.type.has_base(fullnames.ARRAY_FIELD_FULLNAME):
@@ -135,55 +133,99 @@ def process_field_instantiation(ctx: FunctionContext) -> Type:
     return transform_into_proper_return_type(ctx)
 
 
+def _parse_choices_type(ctx: FunctionContext, choices_arg: Expression) -> Optional[str]:
+    if isinstance(choices_arg, (TupleExpr, ListExpr)):
+        # iterable of 2 element tuples of two kinds
+        _, analyzed_choices = ctx.api.analyze_iterable_item_type(choices_arg)
+        if isinstance(analyzed_choices, TupleType):
+            first_element_type = analyzed_choices.items[0]
+            if isinstance(first_element_type, Instance):
+                return first_element_type.type.fullname()
+
+
+def _parse_referenced_model(ctx: FunctionContext, to_arg: Expression) -> Optional[TypeInfo]:
+    if isinstance(to_arg, NameExpr) and isinstance(to_arg.node, TypeInfo):
+        # reference to the model class
+        return to_arg.node
+
+    elif isinstance(to_arg, StrExpr):
+        referenced_model_info = helpers.get_model_info(to_arg.value, ctx.api.modules)
+        if referenced_model_info is not None:
+            return referenced_model_info
+
+
 def parse_field_init_arguments_into_model_metadata(ctx: FunctionContext) -> None:
-    api = cast(TypeChecker, ctx.api)
-    outer_model = api.scope.active_class()
+    outer_model = ctx.api.scope.active_class()
     if outer_model is None or not outer_model.has_base(fullnames.MODEL_CLASS_FULLNAME):
         # outside models.Model class, undetermined
         return
 
-    field_name = None
-    for name_expr, stmt in helpers.iter_over_assignments(outer_model.defn):
-        if stmt == ctx.context and isinstance(name_expr, NameExpr):
-            field_name = name_expr.name
+    # Determine name of the current field
+    for attr_name, stmt in helpers.iter_over_class_level_assignments(outer_model.defn):
+        if stmt == ctx.context:
+            field_name = attr_name
             break
-    if field_name is None:
+    else:
         return
 
-    fields_metadata = metadata.get_fields_metadata(outer_model)
+    model_fields_metadata = metadata.get_fields_metadata(outer_model)
 
     # primary key
     is_primary_key = False
-    primary_key_arg = helpers.get_argument_by_name(ctx, 'primary_key')
+    primary_key_arg = helpers.get_call_argument_by_name(ctx, 'primary_key')
     if primary_key_arg:
         is_primary_key = helpers.parse_bool(primary_key_arg)
-    fields_metadata[field_name] = {'primary_key': is_primary_key}
+    model_fields_metadata[field_name] = {'primary_key': is_primary_key}
 
     # choices
-    choices_arg = helpers.get_argument_by_name(ctx, 'choices')
-    if choices_arg and isinstance(choices_arg, (TupleExpr, ListExpr)):
-        # iterable of 2 element tuples of two kinds
-        _, analyzed_choices = api.analyze_iterable_item_type(choices_arg)
-        if isinstance(analyzed_choices, TupleType):
-            first_element_type = analyzed_choices.items[0]
-            if isinstance(first_element_type, Instance):
-                fields_metadata[field_name]['choices'] = first_element_type.type.fullname()
+    choices_arg = helpers.get_call_argument_by_name(ctx, 'choices')
+    if choices_arg:
+        choices_type_fullname = _parse_choices_type(ctx.api, choices_arg)
+        if choices_type_fullname:
+            model_fields_metadata[field_name]['choices_type'] = choices_type_fullname
 
     # nullability
-    null_arg = helpers.get_argument_by_name(ctx, 'null')
+    null_arg = helpers.get_call_argument_by_name(ctx, 'null')
     is_nullable = False
     if null_arg:
         is_nullable = helpers.parse_bool(null_arg)
-    fields_metadata[field_name]['null'] = is_nullable
+    model_fields_metadata[field_name]['null'] = is_nullable
 
     # is_blankable
-    blank_arg = helpers.get_argument_by_name(ctx, 'blank')
+    blank_arg = helpers.get_call_argument_by_name(ctx, 'blank')
     is_blankable = False
     if blank_arg:
         is_blankable = helpers.parse_bool(blank_arg)
-    fields_metadata[field_name]['blank'] = is_blankable
+    model_fields_metadata[field_name]['blank'] = is_blankable
 
     # default
-    default_arg = helpers.get_argument_by_name(ctx, 'default')
+    default_arg = helpers.get_call_argument_by_name(ctx, 'default')
     if default_arg and not helpers.is_none_expr(default_arg):
-        fields_metadata[field_name]['default_specified'] = True
+        model_fields_metadata[field_name]['default_specified'] = True
+
+    if helpers.has_any_of_bases(ctx.default_return_type.type, fullnames.RELATED_FIELDS_CLASSES):
+        # to
+        to_arg = helpers.get_call_argument_by_name(ctx, 'to')
+        if to_arg:
+            referenced_model = _parse_referenced_model(ctx, to_arg)
+            if referenced_model is not None:
+                model_fields_metadata[field_name]['to'] = referenced_model.fullname()
+            else:
+                model_fields_metadata[field_name]['to'] = to_arg.value
+                # referenced_model = to_arg.value
+                # raise helpers.IncompleteDefnException()
+
+            # model_fields_metadata[field_name]['to'] = referenced_model.fullname()
+            # if referenced_model is not None:
+            #     model_fields_metadata[field_name]['to'] = referenced_model.fullname()
+            # else:
+            #     assert isinstance(to_arg, StrExpr)
+            #     model_fields_metadata[field_name]['to'] = to_arg.value
+
+        # related_name
+        related_name_arg = helpers.get_call_argument_by_name(ctx, 'related_name')
+        if related_name_arg:
+            if isinstance(related_name_arg, StrExpr):
+                model_fields_metadata[field_name]['related_name'] = related_name_arg.value
+            else:
+                model_fields_metadata[field_name]['related_name'] = outer_model.name().lower() + '_set'

mypy_django_plugin/transformers/forms.py

@@ -16,7 +16,7 @@ def extract_proper_type_for_get_form(ctx: MethodContext) -> Type:
     if not isinstance(object_type, Instance):
         return ctx.default_return_type
 
-    form_class_type = helpers.get_argument_type_by_name(ctx, 'form_class')
+    form_class_type = helpers.get_call_argument_type_by_name(ctx, 'form_class')
     if form_class_type is None or isinstance(form_class_type, NoneTyp):
         # extract from specified form_class in metadata
         form_class_fullname = metadata.get_django_metadata(object_type.type).get('form_class', None)

mypy_django_plugin/transformers/models.py

@@ -18,17 +18,20 @@ from mypy_django_plugin.lib import metadata, fullnames, helpers
 class ModelClassInitializer(metaclass=ABCMeta):
     api: SemanticAnalyzerPass2
     model_classdef: ClassDef
+    app_models_mapping: Optional[Dict[str, str]] = None
 
     @classmethod
-    def from_ctx(cls, ctx: ClassDefContext):
-        return cls(api=cast(SemanticAnalyzerPass2, ctx.api), model_classdef=ctx.cls)
+    def from_ctx(cls, ctx: ClassDefContext, app_models_mapping: Optional[Dict[str, str]]):
+        return cls(api=cast(SemanticAnalyzerPass2, ctx.api),
+                   model_classdef=ctx.cls,
+                   app_models_mapping=app_models_mapping)
 
     def get_meta_attribute(self, name: str) -> Optional[Expression]:
         meta_node = helpers.get_nested_meta_node_for_current_class(self.model_classdef.info)
         if meta_node is None:
             return None
 
-        return helpers.get_assigned_value_for_class(meta_node, name)
+        return helpers.get_assignment_stmt_by_name(meta_node, name)
 
     def is_abstract_model(self) -> bool:
         is_abstract_expr = self.get_meta_attribute('abstract')
@@ -46,29 +49,74 @@ class ModelClassInitializer(metaclass=ABCMeta):
         var.is_initialized_in_class = True
         self.model_classdef.info.names[name] = SymbolTableNode(MDEF, var, plugin_generated=True)
 
+    def model_has_name_defined(self, name: str) -> bool:
+        return name in self.model_classdef.info.names
+
     @abstractmethod
     def run(self) -> None:
         raise NotImplementedError()
 
 
 def iter_over_one_to_n_related_fields(klass: ClassDef) -> Iterator[Tuple[NameExpr, CallExpr]]:
-    for lvalue, rvalue in helpers.iter_call_assignments(klass):
-        if (isinstance(lvalue, NameExpr)
-                and isinstance(rvalue.callee, MemberExpr)):
-            if rvalue.callee.fullname in {fullnames.FOREIGN_KEY_FULLNAME,
-                                          fullnames.ONETOONE_FIELD_FULLNAME}:
-                yield lvalue, rvalue
+    for field_name, field_init in helpers.iter_over_field_inits_in_class(klass):
+        field_info = field_init.callee.node
+        assert isinstance(field_info, TypeInfo)
+
+        if helpers.has_any_of_bases(field_init.callee.node, {fullnames.FOREIGN_KEY_FULLNAME,
+                                                             fullnames.ONETOONE_FIELD_FULLNAME}):
+            yield field_name, field_init
 
 
-class SetIdAttrsForRelatedFields(ModelClassInitializer):
+class AddReferencesToRelatedModels(ModelClassInitializer):
+    """
+    For every
+        attr1 = models.ForeignKey(to=MyModel)
+    sets `attr1_id` attribute to the current model.
+    """
+
     def run(self) -> None:
-        for lvalue, rvalue in iter_over_one_to_n_related_fields(self.model_classdef):
-            node_name = lvalue.name + '_id'
-            self.add_new_node_to_model_class(name=node_name,
-                                             typ=self.api.builtin_type('builtins.int'))
+        for field_name, field_init_expr in helpers.iter_over_field_inits_in_class(self.model_classdef):
+            ref_id_name = field_name + '_id'
+            field_info = field_init_expr.callee.node
+            assert isinstance(field_info, TypeInfo)
+
+            if not self.model_has_name_defined(ref_id_name):
+                if helpers.has_any_of_bases(field_info, {fullnames.FOREIGN_KEY_FULLNAME,
+                                                         fullnames.ONETOONE_FIELD_FULLNAME}):
+                    self.add_new_node_to_model_class(name=ref_id_name,
+                                                     typ=self.api.builtin_type('builtins.int'))
+
+        #     field_init_expr.callee.node
+        #
+        # for field_name, field_init_expr in helpers.iter_call_assignments_in_class(self.model_classdef):
+        #     ref_id_name = field_name + '_id'
+        #     if not self.model_has_name_defined(ref_id_name):
+        #         field_class_info = field_init_expr.callee.node
+        #         if not field_class_info:
+        #
+        #         if not field_init_expr.callee.node:
+        #
+        #         if isinstance(field_init_expr.callee.node, TypeInfo) \
+        #                 and helpers.has_any_of_bases(field_init_expr.callee.node,
+        #                                              {fullnames.FOREIGN_KEY_FULLNAME,
+        #                                               fullnames.ONETOONE_FIELD_FULLNAME}):
+        #             self.add_new_node_to_model_class(name=ref_id_name,
+        #                                              typ=self.api.builtin_type('builtins.int'))
 
 
 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:
@@ -77,24 +125,24 @@ class InjectAnyAsBaseForNestedMeta(ModelClassInitializer):
 
 
 class AddDefaultObjectsManager(ModelClassInitializer):
-    def add_new_manager(self, name: str, manager_type: Optional[Instance]) -> None:
+    def _add_new_manager(self, name: str, manager_type: Optional[Instance]) -> None:
         if manager_type is None:
             return None
         self.add_new_node_to_model_class(name, manager_type)
 
-    def add_private_default_manager(self, manager_type: Optional[Instance]) -> None:
+    def _add_private_default_manager(self, manager_type: Optional[Instance]) -> None:
         if manager_type is None:
             return None
         self.add_new_node_to_model_class('_default_manager', manager_type)
 
-    def get_existing_managers(self) -> List[Tuple[str, TypeInfo]]:
+    def _get_existing_managers(self) -> List[Tuple[str, TypeInfo]]:
         managers = []
         for base in self.model_classdef.info.mro:
-            for name_expr, member_expr in helpers.iter_call_assignments(base.defn):
-                manager_name = name_expr.name
-                callee_expr = member_expr.callee
+            for manager_name, call_expr in helpers.iter_call_assignments_in_class(base.defn):
+                callee_expr = call_expr.callee
                 if isinstance(callee_expr, IndexExpr):
                     callee_expr = callee_expr.analyzed.expr
+
                 if isinstance(callee_expr, (MemberExpr, NameExpr)) \
                         and isinstance(callee_expr.node, TypeInfo) \
                         and callee_expr.node.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
@@ -102,12 +150,12 @@ class AddDefaultObjectsManager(ModelClassInitializer):
         return managers
 
     def run(self) -> None:
-        existing_managers = self.get_existing_managers()
+        existing_managers = self._get_existing_managers()
         if existing_managers:
             first_manager_type = None
             for manager_name, manager_type_info in existing_managers:
                 manager_type = Instance(manager_type_info, args=[Instance(self.model_classdef.info, [])])
-                self.add_new_manager(name=manager_name, manager_type=manager_type)
+                self._add_new_manager(name=manager_name, manager_type=manager_type)
                 if first_manager_type is None:
                     first_manager_type = manager_type
         else:
@@ -117,33 +165,46 @@ class AddDefaultObjectsManager(ModelClassInitializer):
 
             first_manager_type = self.api.named_type_or_none(fullnames.MANAGER_CLASS_FULLNAME,
                                                              args=[Instance(self.model_classdef.info, [])])
-            self.add_new_manager('objects', manager_type=first_manager_type)
+            self._add_new_manager('objects', manager_type=first_manager_type)
 
         if self.is_abstract_model():
             return None
         default_manager_name_expr = self.get_meta_attribute('default_manager_name')
         if isinstance(default_manager_name_expr, StrExpr):
-            self.add_private_default_manager(self.model_classdef.info.get(default_manager_name_expr.value).type)
+            self._add_private_default_manager(self.model_classdef.info.get(default_manager_name_expr.value).type)
         else:
-            self.add_private_default_manager(first_manager_type)
+            self._add_private_default_manager(first_manager_type)
 
 
-class AddIdAttributeIfPrimaryKeyTrueIsNotSet(ModelClassInitializer):
+class AddDefaultPrimaryKey(ModelClassInitializer):
+    """
+    Sets default integer `id` attribute, if:
+        * model is not abstract (abstract = False)
+        * there's no field with primary_key=True
+    """
+
     def run(self) -> None:
         if self.is_abstract_model():
-            # no need for .id attr
+            # abstract models cannot be instantiated, and do not need `id` attribute
             return None
 
-        for _, rvalue in helpers.iter_call_assignments(self.model_classdef):
-            if ('primary_key' in rvalue.arg_names
-                    and self.api.parse_bool(rvalue.args[rvalue.arg_names.index('primary_key')])):
+        for _, field_init_expr in helpers.iter_over_field_inits_in_class(self.model_classdef):
+            if ('primary_key' in field_init_expr.arg_names
+                    and self.api.parse_bool(field_init_expr.args[field_init_expr.arg_names.index('primary_key')])):
                 break
         else:
-            self.add_new_node_to_model_class('id', self.api.builtin_type('builtins.object'))
+            self.add_new_node_to_model_class('id', self.api.builtin_type('builtins.int'))
+
+
+def _get_to_expr(field_init_expr) -> Expression:
+    if 'to' in field_init_expr.arg_names:
+        return field_init_expr.args[field_init_expr.arg_names.index('to')]
+    else:
+        return field_init_expr.args[0]
 
 
 class AddRelatedManagers(ModelClassInitializer):
-    def add_related_manager_variable(self, manager_name: str, related_field_type_data: Dict[str, Any]) -> None:
+    def _add_related_manager_variable(self, manager_name: str, related_field_type_data: Dict[str, Any]) -> None:
         # add dummy related manager for use later
         self.add_new_node_to_model_class(manager_name, self.api.builtin_type('builtins.object'))
 
@@ -153,24 +214,41 @@ class AddRelatedManagers(ModelClassInitializer):
 
     def run(self) -> None:
         for module_name, module_file in self.api.modules.items():
-            for model_defn in helpers.iter_over_classdefs(module_file):
-                if not model_defn.info:
-                    self.api.defer()
+            for model_classdef in helpers.iter_over_toplevel_classes(module_file):
+                for field_name, field_init in helpers.iter_over_field_inits_in_class(model_classdef):
+                    field_info = field_init.callee.node
+                    assert isinstance(field_info, TypeInfo)
 
-                for lvalue, field_init in helpers.iter_call_assignments(model_defn):
-                    if is_related_field(field_init, module_file):
-                        try:
-                            referenced_model_fullname = extract_referenced_model_fullname(field_init,
-                                                                                          module_file=module_file,
-                                                                                          all_modules=self.api.modules)
-                        except helpers.SelfReference:
-                            referenced_model_fullname = model_defn.fullname
+                    if helpers.has_any_of_bases(field_info, fullnames.RELATED_FIELDS_CLASSES):
+                        # try:
+                        to_arg_expr = _get_to_expr(field_init)
+                        if isinstance(to_arg_expr, NameExpr):
+                            referenced_model_fullname = module_file.names[to_arg_expr.name].fullname
+                        else:
+                            assert isinstance(to_arg_expr, StrExpr)
+                            value = to_arg_expr.value
+                            if value == 'self':
+                                # reference to the same model class
+                                referenced_model_fullname = model_classdef.fullname
+                            elif '.' not in value:
+                                # reference to class in the current module
+                                referenced_model_fullname = module_name + '.' + value
+                            else:
+                                referenced_model_fullname = self.app_models_mapping[value]
 
-                        except helpers.SameFileModel as exc:
-                            referenced_model_fullname = module_name + '.' + exc.model_cls_name
+                            # referenced_model_fullname = extract_referenced_model_fullname(field_init,
+                            #                                                               module_file=module_file,
+                            #                                                               all_modules=self.api.modules)
+                            # if not referenced_model_fullname:
+                            #     raise helpers.IncompleteDefnException('Cannot parse referenced model fullname')
+
+                        # except helpers.SelfReference:
+                        #     referenced_model_fullname = model_classdef.fullname
+                        #
+                        # except helpers.SameFileModel as exc:
+                        #     referenced_model_fullname = module_name + '.' + exc.model_cls_name
 
                         if self.model_classdef.fullname == referenced_model_fullname:
-                            related_name = model_defn.name.lower() + '_set'
                             if 'related_name' in field_init.arg_names:
                                 related_name_expr = field_init.args[field_init.arg_names.index('related_name')]
                                 if not isinstance(related_name_expr, StrExpr):
@@ -180,9 +258,10 @@ class AddRelatedManagers(ModelClassInitializer):
                                 if related_name == '+':
                                     # No backwards relation is desired
                                     continue
+                            else:
+                                related_name = model_classdef.name.lower() + '_set'
 
                             # Default related_query_name to related_name
-                            related_query_name = related_name
                             if 'related_query_name' in field_init.arg_names:
                                 related_query_name_expr = field_init.args[field_init.arg_names.index('related_query_name')]
                                 if isinstance(related_query_name_expr, StrExpr):
@@ -191,20 +270,24 @@ class AddRelatedManagers(ModelClassInitializer):
                                     # not string 'related_query_name=' is not yet supported
                                     related_query_name = None
                                 # TODO: Handle defaulting to model name if related_name is not set
-
-                            # as long as Model is not a Generic, one level depth is fine
-                            if field_init.callee.name in {'ForeignKey', 'ManyToManyField'}:
-                                field_type_data = {
-                                    'manager': fullnames.RELATED_MANAGER_CLASS_FULLNAME,
-                                    'of': [model_defn.info.fullname()]
-                                }
                             else:
-                                field_type_data = {
-                                    'manager': model_defn.info.fullname(),
-                                    'of': []
-                                }
+                                related_query_name = related_name
 
-                            self.add_related_manager_variable(related_name, related_field_type_data=field_type_data)
+                            # if helpers.has_any_of_bases(field_info, {fullnames.FOREIGN_KEY_FULLNAME,
+                            #                                          fullnames.MANYTOMANY_FIELD_FULLNAME}):
+                            #     # as long as Model is not a Generic, one level depth is fine
+                            #     field_type_data = {
+                            #         'manager': fullnames.RELATED_MANAGER_CLASS_FULLNAME,
+                            #         'of': [model_classdef.info.fullname()]
+                            #     }
+                            # else:
+                            #     field_type_data = {
+                            #         'manager': model_classdef.info.fullname(),
+                            #         'of': []
+                            #     }
+                            self.add_new_node_to_model_class(related_name, self.api.builtin_type('builtins.object'))
+
+                            # self._add_related_manager_variable(related_name, related_field_type_data=field_type_data)
 
                             if related_query_name is not None:
                                 # Only create related_query_name if it is a string literal
@@ -239,22 +322,24 @@ def is_related_field(expr: CallExpr, module_file: MypyFile) -> bool:
     return False
 
 
-def extract_referenced_model_fullname(rvalue_expr: CallExpr,
+def extract_referenced_model_fullname(field_init_expr: CallExpr,
                                       module_file: MypyFile,
                                       all_modules: Dict[str, MypyFile]) -> Optional[str]:
     """ Returns fullname of a Model referenced in "to=" argument of the CallExpr"""
-    if 'to' in rvalue_expr.arg_names:
-        to_expr = rvalue_expr.args[rvalue_expr.arg_names.index('to')]
+    if 'to' in field_init_expr.arg_names:
+        to_expr = field_init_expr.args[field_init_expr.arg_names.index('to')]
     else:
-        to_expr = rvalue_expr.args[0]
+        to_expr = field_init_expr.args[0]
 
     if isinstance(to_expr, NameExpr):
         return module_file.names[to_expr.name].fullname
+
     elif isinstance(to_expr, StrExpr):
         typ_fullname = helpers.get_model_fullname_from_string(to_expr.value, all_modules)
         if typ_fullname is None:
             return None
         return typ_fullname
+
     return None
 
 
@@ -284,16 +369,18 @@ def add_get_set_attr_fallback_to_any(ctx: ClassDefContext):
     add_method(ctx, '__setattr__', [name_arg, value_arg], any)
 
 
-def process_model_class(ctx: ClassDefContext, ignore_unknown_attributes: bool) -> None:
+def process_model_class(ctx: ClassDefContext,
+                        ignore_unknown_attributes: bool,
+                        app_models_mapping: Optional[Dict[str, str]]) -> None:
     initializers = [
         InjectAnyAsBaseForNestedMeta,
+        AddDefaultPrimaryKey,
+        AddReferencesToRelatedModels,
         AddDefaultObjectsManager,
-        AddIdAttributeIfPrimaryKeyTrueIsNotSet,
-        SetIdAttrsForRelatedFields,
         AddRelatedManagers,
     ]
     for initializer_cls in initializers:
-        initializer_cls.from_ctx(ctx).run()
+        initializer_cls.from_ctx(ctx, app_models_mapping).run()
 
     add_dummy_init_method(ctx)
 

mypy_django_plugin/transformers/queryset.py

@@ -110,8 +110,8 @@ def extract_proper_type_queryset_values_list(ctx: MethodContext) -> Type:
             column_names.append(None)
             only_strings_as_fields_expressions = False
 
-    flat = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'flat'))
-    named = helpers.parse_bool(helpers.get_argument_by_name(ctx, 'named'))
+    flat = helpers.parse_bool(helpers.get_call_argument_by_name(ctx, 'flat'))
+    named = helpers.parse_bool(helpers.get_call_argument_by_name(ctx, 'named'))
 
     api = cast(TypeChecker, ctx.api)
     if named and flat:

mypy_django_plugin/transformers/settings.py

@@ -53,7 +53,7 @@ def return_user_model_hook(ctx: FunctionContext, settings_modules: List[str]) ->
     setting_module = api.modules[setting_module_name]
 
     model_path = None
-    for name_expr, rvalue_expr in helpers.iter_over_assignments(setting_module):
+    for name_expr, rvalue_expr in helpers.iter_over_assignments_in_class(setting_module):
         if isinstance(name_expr, NameExpr) and isinstance(rvalue_expr, StrExpr):
             if name_expr.name == 'AUTH_USER_MODEL':
                 model_path = rvalue_expr.value