add support for default related managers, fixes #18

This commit is contained in:
Maxim Kurnikov
2019-02-13 17:11:22 +03:00
parent 70378b8f40
commit 4a22da29cb
2 changed files with 15 additions and 5 deletions

View File

@@ -179,15 +179,17 @@ class AddRelatedManagers(ModelClassInitializer):
return None
if self.model_classdef.fullname == ref_to_fullname:
related_manager_name = defn.name.lower() + '_set'
if 'related_name' in rvalue.arg_names:
related_name_expr = rvalue.args[rvalue.arg_names.index('related_name')]
if not isinstance(related_name_expr, StrExpr):
return None
related_name = related_name_expr.value
related_manager_name = related_name_expr.value
typ = get_related_field_type(rvalue, self.api, defn.info)
if typ is None:
return None
self.add_new_node_to_model_class(related_name, typ)
self.add_new_node_to_model_class(related_manager_name, typ)
def iter_over_classdefs(module_file: MypyFile) -> Iterator[ClassDef]:

View File

@@ -239,3 +239,11 @@ class ParkingSpot(BaseModel):
class Booking(BaseModel):
parking_spot = models.ForeignKey(to=ParkingSpot, null=True, on_delete=models.SET_NULL)
[out]
[CASE if_no_related_name_is_passed_create_default_related_managers]
from django.db import models
class Publisher(models.Model):
pass
class Book(models.Model):
publisher = models.ForeignKey(to=Publisher, on_delete=models.CASCADE)
reveal_type(Publisher().book_set) # E: Revealed type is 'django.db.models.manager.RelatedManager[main.Book]'