Skip to content

Commit

Permalink
Resolve dynamic manager methods through manager MRO (typeddjango#1701)
Browse files Browse the repository at this point in the history
This mainly affects reverse managers as we can't inherit dynamic
managers in any other way.

It also renders a manual code path for method name lookup irrelevant as
we can let mypy do that for us.
  • Loading branch information
flaeppe authored Sep 11, 2023
1 parent c9c3570 commit f5e65d2
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 43 deletions.
43 changes: 4 additions & 39 deletions mypy_django_plugin/transformers/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,10 +58,11 @@ def get_method_type_from_dynamic_manager(
Attempt to resolve a method on a manager that was built from '.from_queryset'
"""

manager_type_info = manager_instance.type
manager_type_info = manager_instance.type.get_containing_type_info(method_name)

if (
"django" not in manager_type_info.metadata
manager_type_info is None
or "django" not in manager_type_info.metadata
or "from_queryset_manager" not in manager_type_info.metadata["django"]
):
# Manager isn't dynamically added
Expand Down Expand Up @@ -235,45 +236,9 @@ def _replace_type_var(ret_type: MypyType, to_replace: str, replace_by: MypyType)
return ret_type


def get_method_type_from_reverse_manager(
api: TypeChecker, method_name: str, manager_type_info: TypeInfo
) -> Optional[ProperType]:
"""
Attempts to resolve a reverse manager's method via the '_default_manager' manager on the related model
From Django docs:
"By default the RelatedManager used for reverse relations is a subclass of the default manager for that model."
Ref: https://docs.djangoproject.com/en/dev/topics/db/queries/#using-a-custom-reverse-manager
"""
is_reverse_manager = (
"django" in manager_type_info.metadata and "related_manager_to_model" in manager_type_info.metadata["django"]
)
if not is_reverse_manager:
return None

related_model_fullname = manager_type_info.metadata["django"]["related_manager_to_model"]
assert isinstance(related_model_fullname, str)
model_info = helpers.lookup_fully_qualified_typeinfo(api, related_model_fullname)
if model_info is None:
return None

# We should _always_ have a '_default_manager' on a model
assert "_default_manager" in model_info.names
assert isinstance(model_info.names["_default_manager"].node, Var)
manager_instance = model_info.names["_default_manager"].node.type
return (
get_method_type_from_dynamic_manager(api, method_name, manager_instance)
# TODO: Can we assert on None and Instance?
if manager_instance is not None and isinstance(manager_instance, Instance)
else None
)


def resolve_manager_method_from_instance(instance: Instance, method_name: str, ctx: AttributeContext) -> MypyType:
api = helpers.get_typechecker_api(ctx)
method_type = get_method_type_from_dynamic_manager(
api, method_name, instance
) or get_method_type_from_reverse_manager(api, method_name, instance.type)

method_type = get_method_type_from_dynamic_manager(api, method_name, instance)
return method_type if method_type is not None else ctx.default_attr_type


Expand Down
12 changes: 8 additions & 4 deletions tests/typecheck/fields/test_related.yml
Original file line number Diff line number Diff line change
Expand Up @@ -642,8 +642,10 @@
reveal_type(user.article_set) # N: Revealed type is "myapp.models.Article_RelatedManager"
reveal_type(user.book_set.add) # N: Revealed type is "def (*objs: Union[myapp.models.Book, builtins.int], *, bulk: builtins.bool =)"
reveal_type(user.article_set.add) # N: Revealed type is "def (*objs: Union[myapp.models.Article, builtins.int], *, bulk: builtins.bool =)"
reveal_type(user.book_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Book]"
reveal_type(user.article_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet[myapp.models.Article]"
reveal_type(user.book_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet"
reveal_type(user.book_set.get()) # N: Revealed type is "myapp.models.Book"
reveal_type(user.article_set.filter) # N: Revealed type is "def (*args: Any, **kwargs: Any) -> myapp.models.LibraryEntityQuerySet"
reveal_type(user.article_set.get()) # N: Revealed type is "myapp.models.Article"
reveal_type(user.book_set.queryset_method()) # N: Revealed type is "builtins.int"
reveal_type(user.article_set.queryset_method()) # N: Revealed type is "builtins.int"
installed_apps:
Expand Down Expand Up @@ -798,11 +800,13 @@
from myapp.models.user import User
reveal_type(Store().purchases) # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager"
reveal_type(Store().purchases.queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
reveal_type(Store().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase]"
reveal_type(Store().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
reveal_type(Store().purchases.get()) # N: Revealed type is "myapp.models.purchase.Purchase"
reveal_type(Store().purchases.filter().queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
reveal_type(User().purchases) # N: Revealed type is "myapp.models.purchase.Purchase_RelatedManager"
reveal_type(User().purchases.queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
reveal_type(User().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet[myapp.models.purchase.Purchase]"
reveal_type(User().purchases.filter()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
reveal_type(User().purchases.get()) # N: Revealed type is "myapp.models.purchase.Purchase"
reveal_type(User().purchases.filter().queryset_method()) # N: Revealed type is "myapp.models.querysets.PurchaseQuerySet"
installed_apps:
- myapp
Expand Down

0 comments on commit f5e65d2

Please sign in to comment.