From 4a22da29cb45be520748dc546a9bf20d01ef719d Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Wed, 13 Feb 2019 17:11:22 +0300
Subject: [PATCH] add support for default related managers, fixes #18

---
 mypy_django_plugin/plugins/models.py    | 12 +++++++-----
 test-data/typecheck/related_fields.test |  8 ++++++++
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/mypy_django_plugin/plugins/models.py b/mypy_django_plugin/plugins/models.py
index abc1bd7..b849c22 100644
--- a/mypy_django_plugin/plugins/models.py
+++ b/mypy_django_plugin/plugins/models.py
@@ -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
-                                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)
+                                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_manager_name, typ)
 
 
 def iter_over_classdefs(module_file: MypyFile) -> Iterator[ClassDef]:
diff --git a/test-data/typecheck/related_fields.test b/test-data/typecheck/related_fields.test
index 4be10e8..6b5dde8 100644
--- a/test-data/typecheck/related_fields.test
+++ b/test-data/typecheck/related_fields.test
@@ -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]'