From ec8717859ef590bbcd66165efd5738f8792cef91 Mon Sep 17 00:00:00 2001 From: Thibaut Decombe Date: Tue, 24 Mar 2026 21:06:00 +0100 Subject: [PATCH] Refactor auto reparametrize to reduce code duplication --- mypy_django_plugin/transformers/managers.py | 96 +++++++++------------ 1 file changed, 39 insertions(+), 57 deletions(-) diff --git a/mypy_django_plugin/transformers/managers.py b/mypy_django_plugin/transformers/managers.py index 8ba746196..0708e19b5 100644 --- a/mypy_django_plugin/transformers/managers.py +++ b/mypy_django_plugin/transformers/managers.py @@ -505,6 +505,11 @@ def populate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeIn def add_as_manager_to_queryset_class(ctx: ClassDefContext) -> None: + """ + Insert a new manager class node for a: ' = .as_manager()'. + + Similar to `create_manager_info_from_from_queryset_call` + """ semanal_api = helpers.get_semanal_api(ctx) def _defer() -> None: @@ -589,44 +594,34 @@ def _defer() -> None: ) -def reparametrize_any_queryset_hook(ctx: ClassDefContext) -> None: +def reparametrize_generic_class(ctx: ClassDefContext, base_class_fullname: str) -> None: """ - Add implicit generics to QuerySet subclasses that are defined without generic. - - Eg. - - class MyQuerySet(models.QuerySet): ... - - is interpreted as: - - _Model = TypeVar('_Model', bound=Model, covariant=True) - _Row = TypeVar('_Row', covariant=True, default=_Model) - class MyQuerySet(models.QuerySet[_Model, _Row]): ... + Add implicit generics to classes that are defined without generic. Note that this does not happen if mypy is run with disallow_any_generics = True, as not specifying the generic type is then considered an error. """ - queryset = ctx.api.lookup_fully_qualified_or_none(ctx.cls.fullname) - if queryset is None or queryset.node is None: + class_info = ctx.api.lookup_fully_qualified_or_none(ctx.cls.fullname) + if class_info is None or class_info.node is None: return - assert isinstance(queryset.node, TypeInfo) + assert isinstance(class_info.node, TypeInfo) - if queryset.node.type_vars: - # We've already been here + if class_info.node.type_vars: + # We've already been here or the class is already declared with generic types return - parent_queryset = next( - (base for base in queryset.node.bases if base.type.has_base(fullnames.QUERYSET_CLASS_FULLNAME)), + parent_class = next( + (base for base in class_info.node.bases if base.type.has_base(base_class_fullname)), None, ) - if parent_queryset is None or len(parent_queryset.args) != 2: + if parent_class is None or len(parent_class.args) < 1: return - model_param = get_proper_type(parent_queryset.args[0]) + model_param = get_proper_type(parent_class.args[0]) if not isinstance(model_param, AnyType) or model_param.type_of_any is not TypeOfAny.from_omitted_generics: return - type_vars = tuple(parent_queryset.type.defn.type_vars) + type_vars = tuple(parent_class.type.defn.type_vars) # If we end up with placeholders we need to defer so the placeholders are # resolved in a future iteration @@ -636,58 +631,45 @@ class MyQuerySet(models.QuerySet[_Model, _Row]): ... else: return - parent_queryset.args = type_vars - queryset.node.defn.type_vars = list(type_vars) - queryset.node.add_type_vars() + parent_class.args = type_vars + class_info.node.defn.type_vars = list(type_vars) + class_info.node.add_type_vars() -def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None: +def reparametrize_any_queryset_hook(ctx: ClassDefContext) -> None: """ - Add implicit generics to manager classes that are defined without generic. + Add implicit generics to QuerySet subclasses that are defined without generic. Eg. - class MyManager(models.Manager): ... + class MyQuerySet(models.QuerySet): ... is interpreted as: - _T = TypeVar('_T', covariant=True) - class MyManager(models.Manager[_T]): ... + _Model = TypeVar('_Model', bound=Model, covariant=True) + _Row = TypeVar('_Row', covariant=True, default=_Model) + class MyQuerySet(models.QuerySet[_Model, _Row]): ... Note that this does not happen if mypy is run with disallow_any_generics = True, as not specifying the generic type is then considered an error. """ + reparametrize_generic_class(ctx, fullnames.QUERYSET_CLASS_FULLNAME) - manager = ctx.api.lookup_fully_qualified_or_none(ctx.cls.fullname) - if manager is None or manager.node is None: - return - assert isinstance(manager.node, TypeInfo) - if manager.node.type_vars: - # We've already been here - return +def reparametrize_any_manager_hook(ctx: ClassDefContext) -> None: + """ + Add implicit generics to manager classes that are defined without generic. - parent_manager = next( - (base for base in manager.node.bases if base.type.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME)), - None, - ) - if parent_manager is None or len(parent_manager.args) != 1: - return + Eg. - model_param = get_proper_type(parent_manager.args[0]) - if not isinstance(model_param, AnyType) or model_param.type_of_any is not TypeOfAny.from_omitted_generics: - return + class MyManager(models.Manager): ... - type_vars = tuple(parent_manager.type.defn.type_vars) + is interpreted as: - # If we end up with placeholders we need to defer so the placeholders are - # resolved in a future iteration - if any(has_placeholder(type_var) for type_var in type_vars): - if not ctx.api.final_iteration: - ctx.api.defer() - else: - return + _T = TypeVar('_T', covariant=True) + class MyManager(models.Manager[_T]): ... - parent_manager.args = type_vars - manager.node.defn.type_vars = list(type_vars) - manager.node.add_type_vars() + Note that this does not happen if mypy is run with disallow_any_generics = True, + as not specifying the generic type is then considered an error. + """ + reparametrize_generic_class(ctx, fullnames.BASE_MANAGER_CLASS_FULLNAME)