mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-14 15:57:08 +08:00
Copy decorated queryset methods to manager too (#646)
* copy decorated queryset methods to manager too * added test for from_manager with decorated queryset methods
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
from mypy.nodes import GDEF, FuncDef, MemberExpr, NameExpr, RefExpr, StrExpr, SymbolTableNode, TypeInfo
|
||||
from mypy.nodes import GDEF, Decorator, FuncDef, MemberExpr, NameExpr, RefExpr, StrExpr, SymbolTableNode, TypeInfo
|
||||
from mypy.plugin import ClassDefContext, DynamicClassDefContext
|
||||
from mypy.types import AnyType, Instance, TypeOfAny
|
||||
|
||||
@@ -67,6 +67,11 @@ def create_new_manager_class_from_from_queryset_method(ctx: DynamicClassDefConte
|
||||
break
|
||||
for name, sym in class_mro_info.names.items():
|
||||
if isinstance(sym.node, FuncDef):
|
||||
func_node = sym.node
|
||||
elif isinstance(sym.node, Decorator):
|
||||
func_node = sym.node.func
|
||||
else:
|
||||
continue
|
||||
helpers.copy_method_to_another_class(
|
||||
class_def_context, self_type, new_method_name=name, method_node=sym.node
|
||||
class_def_context, self_type, new_method_name=name, method_node=func_node
|
||||
)
|
||||
|
||||
@@ -179,3 +179,25 @@
|
||||
class BaseQuerySet(models.QuerySet):
|
||||
def base_queryset_method(self, param: Union[int, str]) -> NoReturn:
|
||||
raise ValueError
|
||||
|
||||
- case: from_queryset_with_decorated_queryset_methods
|
||||
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.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, transaction
|
||||
|
||||
class ModelQuerySet(models.QuerySet):
|
||||
@transaction.atomic
|
||||
def queryset_method(self) -> str:
|
||||
return 'hello'
|
||||
|
||||
NewManager = models.Manager.from_queryset(ModelQuerySet)
|
||||
class MyModel(models.Model):
|
||||
objects = NewManager()
|
||||
|
||||
Reference in New Issue
Block a user