Notify when Manager.from_queryset happens inside model class body (#824)

* Refactor to more easily support additional config options

* Notify when Manager.from_queryset happens inside model class body

- A warning will be emitted whenever `Manager.from_queryset` happens
  inside of a model class body

* Resolve generated default manager types before final iteration

A default manager on a model should always exist, eventually. Although,
we extend to look through dynamically generated managers on each
iteration instead of deferring until the final iteration.

mypy_django_plugin/config.py

@@ -0,0 +1,98 @@
+import configparser
+import textwrap
+from functools import partial
+from pathlib import Path
+from typing import Any, Callable, Dict, NoReturn, Optional
+
+import tomli
+
+INI_USAGE = """
+(config)
+...
+[mypy.plugins.django_stubs]
+    django_settings_module: str (required)
+...
+"""
+TOML_USAGE = """
+(config)
+...
+[tool.django-stubs]
+django_settings_module = str (required)
+...
+"""
+INVALID_FILE = "mypy config file is not specified or found"
+COULD_NOT_LOAD_FILE = "could not load configuration file"
+MISSING_SECTION = "no section [{section}] found".format
+MISSING_DJANGO_SETTINGS = "missing required 'django_settings_module' config"
+INVALID_SETTING = "invalid {key!r}: the setting must be a boolean".format
+
+
+def exit_with_error(msg: str, is_toml: bool = False) -> NoReturn:
+    """Using mypy's argument parser, raise `SystemExit` to fail hard if validation fails.
+
+    Considering that the plugin's startup duration is around double as long as mypy's, this aims to
+    import and construct objects only when that's required - which happens once and terminates the
+    run. Considering that most of the runs are successful, there's no need for this to linger in the
+    global scope.
+    """
+    from mypy.main import CapturableArgumentParser
+
+    handler = CapturableArgumentParser(
+        prog="(django-stubs) mypy", usage=textwrap.dedent(TOML_USAGE if is_toml else INI_USAGE)
+    )
+    handler.error(msg)
+
+
+class DjangoPluginConfig:
+    __slots__ = ("django_settings_module",)
+    django_settings_module: str
+
+    def __init__(self, config_file: Optional[str]) -> None:
+        if not config_file:
+            exit_with_error(INVALID_FILE)
+
+        filepath = Path(config_file)
+        if not filepath.is_file():
+            exit_with_error(INVALID_FILE)
+
+        if filepath.suffix.lower() == ".toml":
+            self.parse_toml_file(filepath)
+        else:
+            self.parse_ini_file(filepath)
+
+    def parse_toml_file(self, filepath: Path) -> None:
+        toml_exit: Callable[[str], NoReturn] = partial(exit_with_error, is_toml=True)
+        try:
+            with filepath.open(mode="rb") as f:
+                data = tomli.load(f)
+        except (tomli.TOMLDecodeError, OSError):
+            toml_exit(COULD_NOT_LOAD_FILE)
+
+        try:
+            config: Dict[str, Any] = data["tool"]["django-stubs"]
+        except KeyError:
+            toml_exit(MISSING_SECTION(section="tool.django-stubs"))
+
+        if "django_settings_module" not in config:
+            toml_exit(MISSING_DJANGO_SETTINGS)
+
+        self.django_settings_module = config["django_settings_module"]
+        if not isinstance(self.django_settings_module, str):
+            toml_exit("invalid 'django_settings_module': the setting must be a string")
+
+    def parse_ini_file(self, filepath: Path) -> None:
+        parser = configparser.ConfigParser()
+        try:
+            with filepath.open(encoding="utf-8") as f:
+                parser.read_file(f, source=str(filepath))
+        except OSError:
+            exit_with_error(COULD_NOT_LOAD_FILE)
+
+        section = "mypy.plugins.django-stubs"
+        if not parser.has_section(section):
+            exit_with_error(MISSING_SECTION(section=section))
+
+        if not parser.has_option(section, "django_settings_module"):
+            exit_with_error(MISSING_DJANGO_SETTINGS)
+
+        self.django_settings_module = parser.get(section, "django_settings_module").strip("'\"")

mypy_django_plugin/errorcodes.py

@@ -0,0 +1,3 @@
+from mypy.errorcodes import ErrorCode
+
+MANAGER_UNTYPED = ErrorCode("django-manager", "Untyped manager disallowed", "Django")

mypy_django_plugin/main.py

@@ -1,10 +1,7 @@
-import configparser
 import sys
-import textwrap
 from functools import partial
-from typing import Callable, Dict, List, NoReturn, Optional, Tuple, cast
+from typing import Callable, Dict, List, Optional, Tuple
 
-import tomli
 from django.db.models.fields.related import RelatedField
 from mypy.modulefinder import mypy_path
 from mypy.nodes import MypyFile, TypeInfo
@@ -21,11 +18,13 @@
 from mypy.types import Type as MypyType
 
 import mypy_django_plugin.transformers.orm_lookups
+from mypy_django_plugin.config import DjangoPluginConfig
 from mypy_django_plugin.django.context import DjangoContext
 from mypy_django_plugin.lib import fullnames, helpers
 from mypy_django_plugin.transformers import fields, forms, init_create, meta, querysets, request, settings
 from mypy_django_plugin.transformers.managers import (
     create_new_manager_class_from_from_queryset_method,
+    fail_if_manager_type_created_in_model_body,
     resolve_manager_method,
 )
 from mypy_django_plugin.transformers.models import (
@@ -60,94 +59,15 @@ def add_new_manager_base_hook(ctx: ClassDefContext) -> None:
     helpers.add_new_manager_base(ctx.api, ctx.cls.fullname)
 
 
-def extract_django_settings_module(config_file_path: Optional[str]) -> str:
-    def exit(error_type: int) -> NoReturn:
-        """Using mypy's argument parser, raise `SystemExit` to fail hard if validation fails.
-
-        Considering that the plugin's startup duration is around double as long as mypy's, this aims to
-        import and construct objects only when that's required - which happens once and terminates the
-        run. Considering that most of the runs are successful, there's no need for this to linger in the
-        global scope.
-        """
-        from mypy.main import CapturableArgumentParser
-
-        usage = """
-        (config)
-        ...
-        [mypy.plugins.django_stubs]
-            django_settings_module: str (required)
-        ...
-        """
-        handler = CapturableArgumentParser(prog="(django-stubs) mypy", usage=textwrap.dedent(usage))
-        messages = {
-            1: "mypy config file is not specified or found",
-            2: "no section [mypy.plugins.django-stubs]",
-            3: "the setting is not provided",
-        }
-        handler.error("'django_settings_module' is not set: " + messages[error_type])
-
-    def exit_toml(error_type: int) -> NoReturn:
-        from mypy.main import CapturableArgumentParser
-
-        usage = """
-        (config)
-        ...
-        [tool.django-stubs]
-        django_settings_module = str (required)
-        ...
-        """
-        handler = CapturableArgumentParser(prog="(django-stubs) mypy", usage=textwrap.dedent(usage))
-        messages = {
-            1: "mypy config file is not specified or found",
-            2: "no section [tool.django-stubs]",
-            3: "the setting is not provided",
-            4: "the setting must be a string",
-        }
-        handler.error("'django_settings_module' not found or invalid: " + messages[error_type])
-
-    if config_file_path and helpers.is_toml(config_file_path):
-        try:
-            with open(config_file_path, encoding="utf-8") as config_file_obj:
-                toml_data = tomli.loads(config_file_obj.read())
-        except Exception:
-            exit_toml(1)
-        try:
-            config = toml_data["tool"]["django-stubs"]
-        except KeyError:
-            exit_toml(2)
-
-        if "django_settings_module" not in config:
-            exit_toml(3)
-
-        if not isinstance(config["django_settings_module"], str):
-            exit_toml(4)
-
-        return config["django_settings_module"]
-    else:
-        parser = configparser.ConfigParser()
-        try:
-            with open(cast(str, config_file_path)) as handle:
-                parser.read_file(handle, source=config_file_path)
-        except (IsADirectoryError, OSError):
-            exit(1)
-
-        section = "mypy.plugins.django-stubs"
-        if not parser.has_section(section):
-            exit(2)
-        settings = parser.get(section, "django_settings_module", fallback=None) or exit(3)
-
-        return settings.strip("'\"")
-
-
 class NewSemanalDjangoPlugin(Plugin):
     def __init__(self, options: Options) -> None:
         super().__init__(options)
-        django_settings_module = extract_django_settings_module(options.config_file)
+        self.plugin_config = DjangoPluginConfig(options.config_file)
         # Add paths from MYPYPATH env var
         sys.path.extend(mypy_path())
         # Add paths from mypy_path config option
         sys.path.extend(options.mypy_path)
-        self.django_context = DjangoContext(django_settings_module)
+        self.django_context = DjangoContext(self.plugin_config.django_settings_module)
 
     def _get_current_queryset_bases(self) -> Dict[str, int]:
         model_sym = self.lookup_fully_qualified(fullnames.QUERYSET_CLASS_FULLNAME)
@@ -306,6 +226,11 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
                 django_context=self.django_context,
             )
 
+        if method_name == "from_queryset":
+            info = self._get_typeinfo_or_none(class_fullname)
+            if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
+                return fail_if_manager_type_created_in_model_body
+
         return None
 
     def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:

mypy_django_plugin/transformers/managers.py

@@ -15,11 +15,12 @@
     TypeInfo,
     Var,
 )
-from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
+from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext
 from mypy.types import AnyType, CallableType, Instance, ProperType
 from mypy.types import Type as MypyType
 from mypy.types import TypeOfAny, TypeVarType, UnboundType, get_proper_type
 
+from mypy_django_plugin import errorcodes
 from mypy_django_plugin.lib import fullnames, helpers
 
 
@@ -278,3 +279,20 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
 
     # Insert the new manager (dynamic) class
     assert semanal_api.add_symbol_table_node(ctx.name, SymbolTableNode(GDEF, new_manager_info, plugin_generated=True))
+
+
+def fail_if_manager_type_created_in_model_body(ctx: MethodContext) -> MypyType:
+    """
+    Method hook that checks if method `<Manager>.from_queryset` is called inside a model class body.
+
+    Doing so won't, for instance, trigger the dynamic class hook(`create_new_manager_class_from_from_queryset_method`)
+    for managers.
+    """
+    api = helpers.get_typechecker_api(ctx)
+    outer_model_info = api.scope.active_class()
+    if not outer_model_info or not outer_model_info.has_base(fullnames.MODEL_CLASS_FULLNAME):
+        # Not inside a model class definition
+        return ctx.default_return_type
+
+    api.fail("`.from_queryset` called from inside model class body", ctx.context, code=errorcodes.MANAGER_UNTYPED)
+    return ctx.default_return_type

mypy_django_plugin/transformers/models.py

@@ -280,17 +280,17 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
         try:
             default_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(default_manager_fullname)
         except helpers.IncompleteDefnException as exc:
-            if not self.api.final_iteration:
-                raise exc
-            else:
-                # On final round, see if the default manager is a generated (dynamic class) manager
-                base_manager_fullname = helpers.get_class_fullname(default_manager_cls.__bases__[0])
-                generated_manager_info = self.get_generated_manager_info(
-                    default_manager_fullname, base_manager_fullname
-                )
-                if generated_manager_info is None:
-                    return
-                default_manager_info = generated_manager_info
+            # Check if default manager could be a generated manager
+            base_manager_fullname = helpers.get_class_fullname(default_manager_cls.__bases__[0])
+            generated_manager_info = self.get_generated_manager_info(default_manager_fullname, base_manager_fullname)
+            if generated_manager_info is None:
+                # Manager doesn't appear to be generated. Unless we're on the final round,
+                # see if another round could help figuring out the default manager type
+                if not self.api.final_iteration:
+                    raise exc
+                else:
+                    return None
+            default_manager_info = generated_manager_info
 
         default_manager = Instance(default_manager_info, [Instance(self.model_classdef.info, [])])
         self.add_new_node_to_model_class("_default_manager", default_manager)
@@ -326,10 +326,8 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
                     related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(
                         fullnames.RELATED_MANAGER_CLASS
                     )  # noqa: E501
-                    # TODO: Use default manager instead of 'objects'
-                    # See: https://docs.djangoproject.com/en/dev/topics/db/queries/#using-a-custom-reverse-manager
-                    objects = related_model_info.get("objects")
-                    if not objects:
+                    default_manager = related_model_info.get("_default_manager")
+                    if not default_manager:
                         raise helpers.IncompleteDefnException()
                 except helpers.IncompleteDefnException as exc:
                     if not self.api.final_iteration:
@@ -339,7 +337,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
 
                 # create new RelatedManager subclass
                 parametrized_related_manager_type = Instance(related_manager_info, [Instance(related_model_info, [])])
-                default_manager_type = objects.type
+                default_manager_type = default_manager.type
                 if default_manager_type is None:
                     default_manager_type = self.try_generate_related_manager(related_model_cls, related_model_info)
                 if (
@@ -360,7 +358,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
     def try_generate_related_manager(
         self, related_model_cls: Type[Model], related_model_info: TypeInfo
     ) -> Optional[Instance]:
-        manager = related_model_cls._meta.managers_map["objects"]
+        manager = related_model_cls._meta.managers_map["_default_manager"]
         base_manager_fullname = helpers.get_class_fullname(manager.__class__.__bases__[0])
         manager_fullname = helpers.get_class_fullname(manager.__class__)
         generated_managers = self.get_generated_manager_mappings(base_manager_fullname)