Skip to content

Commit

Permalink
construct as_manager instance at checker time
Browse files Browse the repository at this point in the history
  • Loading branch information
asottile committed Jul 26, 2024
1 parent 9fa3404 commit 658de86
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 21 deletions.
5 changes: 5 additions & 0 deletions mypy_django_plugin/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
41 changes: 23 additions & 18 deletions mypy_django_plugin/transformers/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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.
Expand Down
6 changes: 3 additions & 3 deletions tests/typecheck/managers/querysets/test_as_manager.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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":
Expand Down Expand Up @@ -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: |
Expand Down Expand Up @@ -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: |
Expand Down

0 comments on commit 658de86

Please sign in to comment.