add support for django.conf.settings.SETTING_NAME

mypy_django_plugin/main.py

@@ -0,0 +1,59 @@
+import os
+from typing import Callable, Optional
+
+from django.conf import Settings
+from mypy.options import Options
+from mypy.plugin import Plugin, FunctionContext, ClassDefContext
+from mypy.types import Type
+
+from mypy_django_plugin import helpers
+from mypy_django_plugin.plugins.objects_queryset import set_objects_queryset_to_model_class
+from mypy_django_plugin.plugins.postgres_fields import determine_type_of_array_field
+from mypy_django_plugin.plugins.related_fields import set_related_name_instance_for_onetoonefield, \
+    set_related_name_manager_for_foreign_key, set_fieldname_attrs_for_related_fields
+from mypy_django_plugin.plugins.setup_settings import DjangoConfSettingsInitializerHook
+
+
+base_model_classes = {helpers.MODEL_CLASS_FULLNAME}
+
+
+def transform_model_class(ctx: ClassDefContext) -> None:
+    base_model_classes.add(ctx.cls.fullname)
+
+    set_fieldname_attrs_for_related_fields(ctx)
+    set_objects_queryset_to_model_class(ctx)
+
+
+class DjangoPlugin(Plugin):
+    def __init__(self,
+                 options: Options) -> None:
+        super().__init__(options)
+        self.django_settings = None
+
+        django_settings_module = os.environ.get('DJANGO_SETTINGS_MODULE')
+        if django_settings_module:
+            self.django_settings = Settings(django_settings_module)
+
+    def get_function_hook(self, fullname: str
+                          ) -> Optional[Callable[[FunctionContext], Type]]:
+        if fullname == helpers.FOREIGN_KEY_FULLNAME:
+            return set_related_name_manager_for_foreign_key
+
+        if fullname == helpers.ONETOONE_FIELD_FULLNAME:
+            return set_related_name_instance_for_onetoonefield
+
+        if fullname == 'django.contrib.postgres.fields.array.ArrayField':
+            return determine_type_of_array_field
+        return None
+
+    def get_base_class_hook(self, fullname: str
+                            ) -> Optional[Callable[[ClassDefContext], None]]:
+        if fullname in base_model_classes:
+            return transform_model_class
+        if fullname == 'django.conf._DjangoConfLazyObject':
+            return DjangoConfSettingsInitializerHook(settings=self.django_settings)
+        return None
+
+
+def plugin(version):
+    return DjangoPlugin