Add inheritance QuerySet support for from_queryset (#275)

* Add testcase for queryset inheritance

* Add PoC

* Add condition for stop to looping over mro

* Change harcoded queryset class name to constant from fullnames
This commit is contained in:
Dima Boger
2019-12-16 20:16:41 +03:00
committed by Maksim Kurnikov
parent fb1593630a
commit b3ed9e4827
2 changed files with 34 additions and 7 deletions

View File

@@ -4,7 +4,7 @@ from mypy.nodes import (
from mypy.plugin import ClassDefContext, DynamicClassDefContext
from mypy.types import AnyType, Instance, TypeOfAny
from mypy_django_plugin.lib import helpers
from mypy_django_plugin.lib import fullnames, helpers
def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefContext) -> None:
@@ -65,9 +65,13 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
class_def_context = ClassDefContext(cls=new_manager_info.defn,
reason=ctx.call, api=semanal_api)
self_type = Instance(new_manager_info, [])
for name, sym in derived_queryset_info.names.items():
if isinstance(sym.node, FuncDef):
helpers.copy_method_to_another_class(class_def_context,
self_type,
new_method_name=name,
method_node=sym.node)
# we need to copy all methods in MRO before django.db.models.query.QuerySet
for class_mro_info in derived_queryset_info.mro:
if class_mro_info.fullname == fullnames.QUERYSET_CLASS_FULLNAME:
break
for name, sym in class_mro_info.names.items():
if isinstance(sym.node, FuncDef):
helpers.copy_method_to_another_class(class_def_context,
self_type,
new_method_name=name,
method_node=sym.node)

View File

@@ -94,3 +94,26 @@
class MyModel(models.Model):
objects = NewManager()
- case: from_queryset_with_class_inheritance
main: |
from myapp.models import MyModel
reveal_type(MyModel().objects) # N: Revealed type is 'myapp.models.MyModel_NewManager[myapp.models.MyModel]'
reveal_type(MyModel().objects.get()) # N: Revealed type is 'myapp.models.MyModel*'
reveal_type(MyModel().objects.queryset_method()) # N: Revealed type is 'builtins.str'
installed_apps:
- myapp
files:
- path: myapp/__init__.py
- path: myapp/models.py
content: |
from django.db import models
from django.db.models.manager import BaseManager
class BaseQuerySet(models.QuerySet):
def queryset_method(self) -> str:
return 'hello'
class ModelQuerySet(BaseQuerySet):
pass
NewManager = BaseManager.from_queryset(ModelQuerySet)
class MyModel(models.Model):
objects = NewManager()