Skip to content

Commit

Permalink
Fix suggested changes from typeddjango#820
Browse files Browse the repository at this point in the history
  • Loading branch information
flaeppe committed Jan 16, 2022
1 parent 99f2838 commit a7cc58d
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 6 deletions.
17 changes: 16 additions & 1 deletion mypy_django_plugin/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -309,37 +309,51 @@ def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], M
return None

def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
# Base class is a Model class definition
if (
fullname in self.django_context.all_registered_model_class_fullnames
or fullname in self._get_current_model_bases()
):
return partial(transform_model_class, django_context=self.django_context)

# Base class is a Manager class definition
if fullname in self._get_current_manager_bases():
return add_new_manager_base_hook

# Base class is a Form class definition
if fullname in self._get_current_form_bases():
return transform_form_class
return None

def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], MypyType]]:
class_name, _, attr_name = fullname.rpartition(".")

# Lookup of a settings variable
if class_name == fullnames.DUMMY_SETTINGS_BASE_CLASS:
return partial(settings.get_type_of_settings_attribute, django_context=self.django_context)

info = self._get_typeinfo_or_none(class_name)

# Lookup of the '.is_superuser' attribute
if info and info.has_base(fullnames.PERMISSION_MIXIN_CLASS_FULLNAME) and attr_name == "is_superuser":
return partial(set_auth_user_model_boolean_fields, django_context=self.django_context)

# Lookup of the 'request.user' attribute
if info and info.has_base(fullnames.HTTPREQUEST_CLASS_FULLNAME) and attr_name == "user":
return partial(request.set_auth_user_model_as_type_for_request_user, django_context=self.django_context)

# Lookup of the 'user.is_staff' or 'user.is_active' attribute
if info and info.has_base(fullnames.ABSTRACT_USER_MODEL_FULLNAME) and attr_name in ("is_staff", "is_active"):
return partial(set_auth_user_model_boolean_fields, django_context=self.django_context)

# Lookup of a method on a dynamically generated manager class
# i.e. a manager class only existing while mypy is running, not collected from the AST
if (
info
and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)
and class_name in self._get_current_manager_bases()
):
return partial(resolve_manager_method, django_context=self.django_context)
return resolve_manager_method

return None

Expand All @@ -352,6 +366,7 @@ def get_type_analyze_hook(self, fullname: str) -> Optional[Callable[[AnalyzeType
return partial(handle_annotated_type, django_context=self.django_context)

def get_dynamic_class_hook(self, fullname: str) -> Optional[Callable[[DynamicClassDefContext], None]]:
# Create a new manager class definition when a manager's '.from_queryset' classmethod is called
if fullname.endswith("from_queryset"):
class_name, _, _ = fullname.rpartition(".")
info = self._get_typeinfo_or_none(class_name)
Expand Down
7 changes: 2 additions & 5 deletions mypy_django_plugin/transformers/managers.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from mypy.types import Type as MypyType
from mypy.types import TypeOfAny, TypeVarType, UnboundType, get_proper_type

from mypy_django_plugin.django.context import DjangoContext
from mypy_django_plugin.lib import fullnames, helpers


Expand Down Expand Up @@ -96,7 +95,7 @@ def get_method_type_from_reverse_manager(
)


def resolve_manager_method(ctx: AttributeContext, django_context: DjangoContext) -> MypyType:
def resolve_manager_method(ctx: AttributeContext) -> MypyType:
"""
A 'get_attribute_hook' that is intended to be invoked whenever the TypeChecker encounters
an attribute on a class that has 'django.db.models.BaseManager' as a base.
Expand Down Expand Up @@ -179,9 +178,7 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
new_manager_info.metaclass_type = new_manager_info.calculate_metaclass_type()
# Stash the queryset fullname which was passed to .from_queryset
# So that our 'resolve_manager_method' attribute hook can fetch the method from that QuerySet class
new_manager_info.metadata.setdefault("django", {})
new_manager_info.metadata["django"].setdefault("from_queryset_manager", {})
new_manager_info.metadata["django"]["from_queryset_manager"] = derived_queryset_fullname
new_manager_info.metadata["django"] = {"from_queryset_manager": derived_queryset_fullname}

if len(ctx.call.args) > 1:
expr = ctx.call.args[1]
Expand Down
1 change: 1 addition & 0 deletions mypy_django_plugin/transformers/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,7 @@ def run_with_model_cls(self, model_cls: Type[Model]) -> None:
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
Expand Down

0 comments on commit a7cc58d

Please sign in to comment.