Add TypeVar resolution for as_manager querysets (typeddjango#1646)

[moranabadie]

I have made things!

In this pull request, I introduce a new utility function called _find_type_var, which aims to identify the _CTE associated with _MyModelQuerySet. This change is demonstrated in the accompanying test.

I've strived to make the implementation as generic as possible. I believe this approach should not introduce any regressions, although I welcome feedback on whether there might be a simpler way to achieve the same result.

Please refer to the included test for a detailed demonstration of the issue that this pull request seeks to resolve.

Related issues

• Partially Closes create_new_manager_class_from_as_manager_method : TypeVar QuerySet does not seems to work. #1646

[sobolevn]

ALternative idea: how har it would be to resolve these type vars when we create a type with .as_manager() and .as_queryset()?

[moranabadie]

ALternative idea: how har it would be to resolve these type vars when we create a type with .as_manager() and .as_queryset()?

If I understand well you want this to make it work for from_queryset as well. It seems to work (I have added a test).

[sobolevn]

Nice that it works, but I was asking about these lines:

django-stubs/mypy_django_plugin/main.py

Lines 326 to 333 in a8e42cb

	if method_name == "from_queryset":
	info = self._get_typeinfo_or_none(class_name)
	if info and info.has_base(fullnames.BASE_MANAGER_CLASS_FULLNAME):
	return create_new_manager_class_from_from_queryset_method
	elif method_name == "as_manager":
	info = self._get_typeinfo_or_none(class_name)
	if info and info.has_base(fullnames.QUERYSET_CLASS_FULLNAME):
	return create_new_manager_class_from_as_manager_method

Can we do it somewhere there? Or will it be too hard?

[moranabadie]

It seems that all the methods of the generated manager are all purposely implemented with the sym_type=AnyType(TypeOfAny.implementation_artifact), according the comment :

https://github.com/typeddjango/django-stubs/blob/a8e42cbcca7a2ce3d890789223d3043412e51c4c/mypy_django_plugin/transformers/managers.py#L343.py#L352-L362.

populate_manager_from_queryset is called by create_manager_info_from_from_queryset_call and create_new_manager_class_from_as_manager_method.

This I am not sure if it is a good idea to make a special case for the issue of this pull request. Maybe I'm wrong.

[sobolevn]

Ok then, but I would like to have a second opinion as well :)

[moranabadie]

I can also add a more generic approach :

    if isinstance(ret_type, TypeVarType):
        ret_type = _find_type_var(queryset_info, ret_type) or ret_type
    elif isinstance(ret_type, Instance):
        ret_type.args = tuple(
            _find_type_var(queryset_info, item) or item
            if isinstance(item, TypeVarType) else item for item in ret_type.args
        )
    elif isinstance(ret_type, TypeType) and isinstance(ret_type.item, TypeVarType):
        ret_type.item = _find_type_var(queryset_info, ret_type.item) or ret_type.item
    elif isinstance(ret_type, ProperType) and hasattr(ret_type, "items"):
        ret_type.items = [
            _find_type_var(queryset_info, item) or item
            if isinstance(item, TypeVarType) else item for item in ret_type.items
        ]

for

    def example2(self) -> list[_CTE]: ...


    def example3(self) -> type[_CTE]: ...

    def example4(self) -> tuple[_CTE]: ...

    def example5(self) -> List[_CTE]: ...

if you think this is useful and not overkill.

[flaeppe]

I think I found a failing case but hopefully an approach that should fix it.

[flaeppe]

I also realised we have to resolve any TypeVar arguments of the methods. As it's the same case for those as for the return type.

I'm getting a feeling there should be some other way to get the resolved types of a type var. Not at all sure how mypy does it, I'm hoping there could be something to piggyback from there..

[moranabadie]

I also realised we have to resolve any TypeVar arguments of the methods. As it's the same case for those as for the return type.

I'm getting a feeling there should be some other way to get the resolved types of a type var. Not at all sure how mypy does it, I'm hoping there could be something to piggyback from there..

Do you know why those things (setting the return type of the methods) are not set here https://github.com/typeddjango/django-stubs/blob/master/mypy_django_plugin/transformers/managers.py#L351, like @sobolevn said ?

I think it would be easier to use mypy here.
I have also no clue how mypy does it. But I do not think intuitively mypy make this kind of thing by managing the methods and the class separately like it is done in this mypy plugin (populate_manager_from_queryset for creating the manager, and get_method_type_from_dynamic_manager for the methods).

[flaeppe]

Hm, a different way to look at it could be that since;

• We (should) require the original Manager and QuerySet arguments to align
• We create a subclass manager instance from those 2 classes
• We get the manager instance from the attribute hook context

We could just populate any type vars with manager_instance.args[0](?). So what I think I'm saying is that we should fall back on everything else during creation of the new manager type to be correct and just put that in? At least for the model argument..

[flaeppe]

Do you know why those things (setting the return type of the methods) are not set here

Yes, it's because that's happening during semantic analysis and the attribute hook is coming in during type checking. Main practical difference (in my very own words) is that not all types is or probably will be resolved during semantic analysis. So we wait for the type checker to have come around, as we then should have more types to look up and work with.

[...] But I do not think intuitively mypy make this kind of thing by managing the methods and the class separately like it is done in this mypy plugin (populate_manager_from_queryset for creating the manager, and get_method_type_from_dynamic_manager for the methods).

You can backtrack to my (very old) PR as to why we do this: #820 . What Django is doing is that it dynamically creates a new type, based on the manager and queryset class. And we need to help mypy understand that happened and how.

[moranabadie]

Hm, a different way to look at it could be that since;

* We (should) require the original `Manager` and `QuerySet` arguments to align

* We create a subclass manager instance from those 2 classes

* We get the manager _instance_ from the attribute hook context

We could just populate any type vars with manager_instance.args[0](?). So what I think I'm saying is that we should fall back on everything else during creation of the new manager type to be correct and just put that in? At least for the model argument..

I am not sure to follow :

with this example :

from typing import TypeVar
from django.db import models

_CTE = TypeVar("_CTE", bound=models.Model)
_CTE_2 = TypeVar("_CTE_2", bound=models.Model)
_CTE_3 = TypeVar("_CTE_3")

class _MyModelQuerySet(models.QuerySet[_CTE], Generic[_CTE, _CTE_3):
    def example(self) -> _CTE: ...
    def example3(self) -> _CTE_3: ...
class _MyModelQuerySet2(_MyModelQuerySet[_CTE_2, str]):
    def example_2(self) -> _CTE_2: ...

class MyModelQuerySet(_MyModelQuerySet2["MyModel"]): ....

class MyModel(models.Model):
    objects = MyModelQuerySet.as_manager()

in the actual context of get_method_type_from_dynamic_manager, we have def example(self) -> _CTE: .... How can we know that _CTE is related to manager_instance.args[0] ?
For example def example3(self) -> _CTE_3: ... has nothing to do with manager_instance.args[0], i.e. "MyModel"].

Or maybe I did not understand want you want to do :)

[moranabadie]

I also realised we have to resolve any TypeVar arguments of the methods. As it's the same case for those as for the return type.

I added the code and a test to make it work whit the original way.

But indeed it is seems like recreating some mypy code. But with this architecture of django-stubs, I am not sure how to make it without "duplicate code"

[flaeppe]

in the actual context of get_method_type_from_dynamic_manager, we have def example(self) -> _CTE: .... How can we know that _CTE is related to manager_instance.args[0] ?
For example def example3(self) -> _CTE_3: ... has nothing to do with manager_instance.args[0], i.e. "MyModel"].

I'll try to explain what I'm thinking as best as I can. It's rather lengthy though..

So, the first part is to look at what the runtime does. In short words; it creates a new Manager type, which subclasses e.g. django.db.models.Manager and also "copies over" methods from the given queryset. As such, the resulting class is a sort of fusion between the specified manager and queryset.

e.g.

class MyQuerySet(models.QuerySet):
    def method(self):
        ...

class MyManager(models.Manager):
    ...

MyModelManager = MyManager.from_queryset(MyQuerySet)
class MyModel(models.Model):
    objects = MyModelManager()

MyModelManager becomes

class MyModelManager(MyManager):
    method = MyQuerySet.method

The plugin mimics this behaviour.

The second part is to look at the definition of models.BaseManager and models.QuerySet in the stubs.

BaseManager:

django-stubs/django-stubs/db/models/manager.pyi

Lines 12 to 14 in 3a1f139

	_T = TypeVar("_T", bound=Model, covariant=True)
	class BaseManager(Generic[_T]):

QuerySet:

django-stubs/django-stubs/db/models/query.pyi

Lines 11 to 12 in 3a1f139

	_T = TypeVar("_T", bound=Model, covariant=True)
	_Row = TypeVar("_Row", covariant=True)

django-stubs/django-stubs/db/models/query.pyi

Line 43 in 3a1f139

class _QuerySet(Generic[_T, _Row], Collection[_Row], Reversible[_Row], Sized):

django-stubs/django-stubs/db/models/query.pyi

Line 219 in 3a1f139

QuerySet: TypeAlias = _QuerySet[_T, _T]

QuerySet and BaseManager are both generic over a form of models.Model. As a subclass of a BaseManager, our generated manager class (MyModelManager) is also generic over a form of models.Model.

Third part is that when a manager instance is assigned to an attribute of a models.Model, the generic argument of our manager is populated(i.e. this is where we get MyModelManager[MyModel]). This is done by the plugin. Also note that this is where we get the manager instance from and the model argument type.

django-stubs/mypy_django_plugin/transformers/models.py

Lines 320 to 328 in 3a1f139

	for manager_name, manager in model_cls._meta.managers_map.items():
	manager_node = self.model_classdef.info.names.get(manager_name, None)
	manager_fullname = helpers.get_class_fullname(manager.__class__)
	manager_info = self.lookup_manager(manager_fullname, manager)
	if manager_node and manager_node.type is not None:
	# Manager is already typed -> do nothing unless it's a dynamically generated manager
	self.reparametrize_dynamically_created_manager(manager_name, manager_info)
	continue

Now, if our custom def method on our queryset is annotated to return a form of model e.g.

T = TypeVar("T", bound=models.Model)
class MyQuerySet(models.QuerySet[T]):
    def method(self) -> T:
        ...

At runtime, when is the first time we can say anything about model T? -> It's not until our manager has been instantiated and added as an attribute to a model (which is happening with a setattr found here. The call chain to get here is quite lengthy and circular so I'll just link this one place) and also gotten its model attribute (you can see the self.model = cls).

Essentially our plugin runs in to the same situation. Now, with the copies of the queryset methods on our dynamically created manager(MyModelManager), what should T be? -> It has to come from the manager instance, they will always align runtime. To declare that statically we need compatible manager and queryset model arguments. It's a prerequisite to get things right and we should introduce a check that controls that compatibility.

Otherwise the provided manager and queryset classes are incompatible.

But once we reach the code that should resolve/populate T, that code should just replace T with the model argument of the manager instance, as that what is happening runtime. There is no queryset instance.

The check mentioned above should uphold the prerequisite and our code inside get_method_type_from_dynamic_manager can instead be naive and just use the model argument, whatever type it has.

What do you think? Does this make sense?

---

An additional thing with this is if one decides to add more generic arguments to a subclass of a QuerySet(like you gave an example of). But I think we could do stuff in a step by step manner and see what needs to be done once that becomes a real use case.

[flaeppe]

Nice work! I had a few more comments

[flaeppe]

Can we move this check to where the manager and queryset are "fused" together?

django-stubs/mypy_django_plugin/transformers/managers.py

Lines 326 to 329 in 218d5fd

	def populate_manager_from_queryset(manager_info: TypeInfo, queryset_info: TypeInfo) -> None:
	"""
	Add methods from the QuerySet class to the manager.
	"""

That would be more efficient, since it'll only happen once, while this is called for each accessed manager/queryset method. We would also get better context if we'd emit an error message. Then we can only do replacement of type vars in here.

We should also ensure that upper bounds of any TypeVars are compatible (i.e. mypy.nodes.TypeVarExpr.upper_bound or mypy.types.TypeVarType.upper_bound)

[flaeppe]

Hm, thinking twice about this, it's quite branched off and self contained. It's also only new additions and shouldn't really break anything. While it solves a couple of simple use cases.

We could try to improve in subsequent PR:s instead, if you'd prefer that? But I'd say we at least remove _get_base_containing_method in favor of mypy's implementation.

[moranabadie]

yes maybe let's do that in another PR.

[flaeppe]

Nice work 👍 Lets merge this

[flaeppe]

Nice work 👍 Lets merge this

Thought CI had run, seems to be a few syntax fixes for older python versions to fix first 🙂

[moranabadie]

Nice work 👍 Lets merge this

Thought CI had run, seems to be a few syntax fixes for older python versions to fix first 🙂

Arf, I forgot about the | for py<3.10, fixed :)

[moranabadie]

I do not understand why I do have this error https://github.com/typeddjango/django-stubs/actions/runs/6077373631/job/16486992087?pr=1666. I added from __future__ import annotations but I do not understand why I have it only with 3.8 and not 3.7 for example

[flaeppe]

I do not understand why I do have this error https://github.com/typeddjango/django-stubs/actions/runs/6077373631/job/16486992087?pr=1666. I added from __future__ import annotations but I do not understand why I have it only with 3.8 and not 3.7 for example

I'm not sure why actually