mirror of
https://github.com/davidhalter/django-stubs.git
synced 2025-12-06 12:14:28 +08:00
Make related manager inherit from objects of related model (#278)
* related manager inherits from objects of related model * fix test typechecking * lint
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
from typing import Dict, Optional, Type, cast
|
||||
from typing import Dict, List, Optional, Type, cast
|
||||
|
||||
from django.db.models.base import Model
|
||||
from django.db.models.fields import DateField, DateTimeField
|
||||
@@ -58,6 +58,12 @@ class ModelClassInitializer:
|
||||
name=name,
|
||||
sym_type=typ)
|
||||
|
||||
def add_new_class_for_current_module(self, name: str, bases: List[Instance]) -> TypeInfo:
|
||||
current_module = self.api.modules[self.model_classdef.info.module_name]
|
||||
new_class_info = helpers.add_new_class_for_module(current_module,
|
||||
name=name, bases=bases)
|
||||
return new_class_info
|
||||
|
||||
def run(self) -> None:
|
||||
model_cls = self.django_context.get_model_class_by_fullname(self.model_classdef.fullname)
|
||||
if model_cls is None:
|
||||
@@ -164,14 +170,12 @@ class AddManagers(ModelClassInitializer):
|
||||
[Instance(self.model_classdef.info, [])])
|
||||
bases.append(original_base)
|
||||
|
||||
current_module = self.api.modules[self.model_classdef.info.module_name]
|
||||
custom_manager_info = helpers.add_new_class_for_module(current_module,
|
||||
name=name, bases=bases)
|
||||
new_manager_info = self.add_new_class_for_current_module(name, bases)
|
||||
# copy fields to a new manager
|
||||
new_cls_def_context = ClassDefContext(cls=custom_manager_info.defn,
|
||||
new_cls_def_context = ClassDefContext(cls=new_manager_info.defn,
|
||||
reason=self.ctx.reason,
|
||||
api=self.api)
|
||||
custom_manager_type = Instance(custom_manager_info, [Instance(self.model_classdef.info, [])])
|
||||
custom_manager_type = Instance(new_manager_info, [Instance(self.model_classdef.info, [])])
|
||||
|
||||
for name, sym in base_manager_info.names.items():
|
||||
# replace self type with new class, if copying method
|
||||
@@ -185,10 +189,10 @@ class AddManagers(ModelClassInitializer):
|
||||
new_sym = sym.copy()
|
||||
if isinstance(new_sym.node, Var):
|
||||
new_var = Var(name, type=sym.type)
|
||||
new_var.info = custom_manager_info
|
||||
new_var._fullname = custom_manager_info.fullname + '.' + name
|
||||
new_var.info = new_manager_info
|
||||
new_var._fullname = new_manager_info.fullname + '.' + name
|
||||
new_sym.node = new_var
|
||||
custom_manager_info.names[name] = new_sym
|
||||
new_manager_info.names[name] = new_sym
|
||||
|
||||
return custom_manager_type
|
||||
|
||||
@@ -268,15 +272,30 @@ class AddRelatedManagers(ModelClassInitializer):
|
||||
|
||||
if isinstance(relation, (ManyToOneRel, ManyToManyRel)):
|
||||
try:
|
||||
manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS)
|
||||
related_manager_info = self.lookup_typeinfo_or_incomplete_defn_error(fullnames.RELATED_MANAGER_CLASS) # noqa: E501
|
||||
if 'objects' not in related_model_info.names:
|
||||
raise helpers.IncompleteDefnException()
|
||||
except helpers.IncompleteDefnException as exc:
|
||||
if not self.api.final_iteration:
|
||||
raise exc
|
||||
else:
|
||||
continue
|
||||
self.add_new_node_to_model_class(attname,
|
||||
Instance(manager_info, [Instance(related_model_info, [])]))
|
||||
continue
|
||||
|
||||
# create new RelatedManager subclass
|
||||
parametrized_related_manager_type = Instance(related_manager_info,
|
||||
[Instance(related_model_info, [])])
|
||||
default_manager_type = related_model_info.names['objects'].type
|
||||
if (default_manager_type is None
|
||||
or not isinstance(default_manager_type, Instance)
|
||||
or default_manager_type.type.fullname == fullnames.MANAGER_CLASS_FULLNAME):
|
||||
self.add_new_node_to_model_class(attname, parametrized_related_manager_type)
|
||||
continue
|
||||
|
||||
name = related_model_cls.__name__ + '_' + 'RelatedManager'
|
||||
bases = [parametrized_related_manager_type, default_manager_type]
|
||||
new_related_manager_info = self.add_new_class_for_current_module(name, bases)
|
||||
|
||||
self.add_new_node_to_model_class(attname, Instance(new_related_manager_info, []))
|
||||
|
||||
|
||||
class AddExtraFieldMethods(ModelClassInitializer):
|
||||
|
||||
@@ -248,7 +248,7 @@ IGNORED_ERRORS = {
|
||||
re.compile(r'"(HttpRequest|WSGIRequest)" has no attribute'),
|
||||
],
|
||||
'many_to_many': [
|
||||
'(expression has type "List[Article]", variable has type "RelatedManager[Article]"',
|
||||
'(expression has type "List[Article]", variable has type "Article_RelatedManager2',
|
||||
'"add" of "RelatedManager" has incompatible type "Article"; expected "Union[Publication, int]"',
|
||||
],
|
||||
'many_to_one': [
|
||||
|
||||
@@ -648,3 +648,27 @@
|
||||
abstract = True
|
||||
class User(AbstractUser):
|
||||
pass
|
||||
|
||||
|
||||
- case: related_manager_is_a_subclass_of_default_manager
|
||||
main: |
|
||||
from myapp.models import User
|
||||
reveal_type(User().orders) # N: Revealed type is 'myapp.models.Order_RelatedManager'
|
||||
reveal_type(User().orders.get()) # N: Revealed type is 'myapp.models.Order*'
|
||||
reveal_type(User().orders.manager_method()) # N: Revealed type is 'builtins.int'
|
||||
installed_apps:
|
||||
- myapp
|
||||
files:
|
||||
- path: myapp/__init__.py
|
||||
- path: myapp/models.py
|
||||
content: |
|
||||
from django.db import models
|
||||
class User(models.Model):
|
||||
pass
|
||||
class OrderManager(models.Manager):
|
||||
def manager_method(self) -> int:
|
||||
pass
|
||||
class Order(models.Model):
|
||||
objects = OrderManager()
|
||||
user = models.ForeignKey(to=User, on_delete=models.CASCADE, related_name='orders')
|
||||
|
||||
|
||||
Reference in New Issue
Block a user