From 658de866ccb246d0cdac364d5ba226ca1115fd87 Mon Sep 17 00:00:00 2001
From: Anthony Sottile <asottile@umich.edu>
Date: Fri, 26 Jul 2024 13:26:18 -0400
Subject: [PATCH] construct as_manager instance at checker time
---
mypy_django_plugin/main.py | 5 +++
mypy_django_plugin/transformers/managers.py | 41 +++++++++++--------
.../managers/querysets/test_as_manager.yml | 6 +--
3 files changed, 31 insertions(+), 21 deletions(-)
diff --git a/mypy_django_plugin/main.py b/mypy_django_plugin/main.py
index e9f549b70..554550378 100644
--- a/mypy_django_plugin/main.py
+++ b/mypy_django_plugin/main.py
@@ -37,6 +37,7 @@
)
from mypy_django_plugin.transformers.functional import resolve_str_promise_attribute
from mypy_django_plugin.transformers.managers import (
+ construct_as_manager_instance,
create_new_manager_class_from_as_manager_method,
create_new_manager_class_from_from_queryset_method,
reparametrize_any_manager_hook,
@@ -208,6 +209,10 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
fullnames.REVERSE_MANY_TO_ONE_DESCRIPTOR: manytoone.refine_many_to_one_related_manager,
}
return hooks.get(class_fullname)
+ elif method_name == "as_manager":
+ info = self._get_typeinfo_or_none(class_fullname)
+ if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
+ return partial(construct_as_manager_instance, info=info)
if method_name in self.manager_and_queryset_method_hooks:
info = self._get_typeinfo_or_none(class_fullname)
diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py
index 79f8c71cb..ec13bf20c 100644
--- a/mypy_django_plugin/transformers/managers.py
+++ b/mypy_django_plugin/transformers/managers.py
@@ -15,9 +15,8 @@
StrExpr,
SymbolTableNode,
TypeInfo,
- Var,
)
-from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext
+from mypy.plugin import AttributeContext, ClassDefContext, DynamicClassDefContext, MethodContext
from mypy.semanal import SemanticAnalyzer
from mypy.semanal_shared import has_placeholder
from mypy.subtypes import find_member
@@ -552,23 +551,9 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
manager_name=manager_class_name,
manager_base=manager_base,
)
+ queryset_info.metadata.setdefault("django_as_manager_names", {})
+ queryset_info.metadata["django_as_manager_names"][semanal_api.cur_mod_id] = new_manager_info.name
- # Whenever `<QuerySet>.as_manager()` isn't called at class level, we want to ensure
- # that the variable is an instance of our generated manager. Instead of the return
- # value of `.as_manager()`. Though model argument is populated as `Any`.
- # `transformers.models.AddManagers` will populate a model's manager(s), when it
- # finds it on class level.
- var = Var(name=ctx.name, type=Instance(new_manager_info, [AnyType(TypeOfAny.from_omitted_generics)]))
- var.info = new_manager_info
- var._fullname = f"{current_module.fullname}.{ctx.name}"
- var.is_inferred = True
- # Note: Order of `add_symbol_table_node` calls matters. Depending on what level
- # we've found the `.as_manager()` call. Point here being that we want to replace the
- # `.as_manager` return value with our newly created manager.
- added = semanal_api.add_symbol_table_node(
- ctx.name, SymbolTableNode(semanal_api.current_symbol_kind(), var, plugin_generated=True)
- )
- assert added
# Add the new manager to the current module
added = semanal_api.add_symbol_table_node(
# We'll use `new_manager_info.name` instead of `manager_class_name` here
@@ -580,6 +565,26 @@ def create_new_manager_class_from_as_manager_method(ctx: DynamicClassDefContext)
assert added
+def construct_as_manager_instance(ctx: MethodContext, *, info: TypeInfo) -> MypyType:
+ api = helpers.get_typechecker_api(ctx)
+ module = helpers.get_current_module(api)
+ try:
+ manager_name = info.metadata["django_as_manager_names"][module.fullname]
+ except KeyError:
+ return ctx.default_return_type
+
+ manager_node = api.lookup(manager_name)
+ if not isinstance(manager_node.node, TypeInfo):
+ return ctx.default_return_type
+
+ # Whenever `<QuerySet>.as_manager()` isn't called at class level, we want to ensure
+ # that the variable is an instance of our generated manager. Instead of the return
+ # value of `.as_manager()`. Though model argument is populated as `Any`.
+ # `transformers.models.AddManagers` will populate a model's manager(s), when it
+ # finds it on class level.
+ return Instance(manager_node.node, [AnyType(TypeOfAny.from_omitted_generics)])
+
+
def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None:
"""
Add implicit generics to manager classes that are defined without generic.
diff --git a/tests/typecheck/managers/querysets/test_as_manager.yml b/tests/typecheck/managers/querysets/test_as_manager.yml
index 9bdfec735..641d4be30 100644
--- a/tests/typecheck/managers/querysets/test_as_manager.yml
+++ b/tests/typecheck/managers/querysets/test_as_manager.yml
@@ -28,7 +28,7 @@
def just_int(self) -> int: ...
class MyModel(models.Model):
- objects = MyQuerySet.as_manager() # type: ignore[var-annotated]
+ objects = MyQuerySet.as_manager()
class QuerySetWithoutSelf(models.QuerySet["MyModelWithoutSelf"]):
def method(self) -> "QuerySetWithoutSelf":
@@ -75,7 +75,7 @@
...
class MyModel(models.Model):
- objects = MyQuerySet.as_manager() # type: ignore[var-annotated]
+ objects = MyQuerySet.as_manager()
- case: model_gets_generated_manager_as_default_manager
main: |
@@ -285,7 +285,7 @@
objects = MyModelQuerySet.as_manager()
class MyOtherModel(models.Model):
- objects = _MyModelQuerySet2.as_manager() # type: ignore
+ objects = _MyModelQuerySet2.as_manager()
- case: handles_type_vars
main: |