Refactor auto reparametrize to reduce code duplication

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: '<Manager> = <QuerySet>.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)