From c32a7842d6d0a10c0709a195151ee4df4fb1bbb4 Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Sat, 10 Nov 2018 20:03:09 +0300
Subject: [PATCH] add ForeignKey

---
 django-stubs/db/models/__init__.pyi       |  2 ++
 django-stubs/db/models/deletion.pyi       |  2 ++
 django-stubs/db/models/fields/related.pyi | 14 ++++++++++++++
 test/test-data/check-model-relations.test | 14 ++++++++++++++
 test/testdjango.py                        |  3 ++-
 5 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 django-stubs/db/models/deletion.pyi
 create mode 100644 django-stubs/db/models/fields/related.pyi
 create mode 100644 test/test-data/check-model-relations.test

diff --git a/django-stubs/db/models/__init__.pyi b/django-stubs/db/models/__init__.pyi
index ccafd1f..d167319 100644
--- a/django-stubs/db/models/__init__.pyi
+++ b/django-stubs/db/models/__init__.pyi
@@ -7,3 +7,5 @@ from .fields import (AutoField as AutoField,
                      Field as Field,
                      SlugField as SlugField,
                      TextField as TextField)
+from .fields.related import (ForeignKey as ForeignKey)
+from .deletion import CASCADE as CASCADE
diff --git a/django-stubs/db/models/deletion.pyi b/django-stubs/db/models/deletion.pyi
new file mode 100644
index 0000000..45dc807
--- /dev/null
+++ b/django-stubs/db/models/deletion.pyi
@@ -0,0 +1,2 @@
+def CASCADE(collector, field, sub_objs, using):
+    ...
diff --git a/django-stubs/db/models/fields/related.pyi b/django-stubs/db/models/fields/related.pyi
new file mode 100644
index 0000000..763e7a7
--- /dev/null
+++ b/django-stubs/db/models/fields/related.pyi
@@ -0,0 +1,14 @@
+from typing import Type, Union, TypeVar, Any, Generic
+
+from django.db import models
+from django.db.models import Field
+
+_T = TypeVar('_T', bound=models.Model)
+
+
+class ForeignKey(Field, Generic[_T]):
+    def __init__(self,
+                 to: Union[Type[_T], str],
+                 on_delete: Any,
+                 **kwargs): ...
+    def __get__(self, instance, owner) -> _T: ...
diff --git a/test/test-data/check-model-relations.test b/test/test-data/check-model-relations.test
new file mode 100644
index 0000000..33ca653
--- /dev/null
+++ b/test/test-data/check-model-relations.test
@@ -0,0 +1,14 @@
+[case testForeignKeyWithClass]
+from django.db import models
+
+class User(models.Model):
+    pass
+
+class Profile(models.Model):
+    user = models.ForeignKey(to=User, on_delete=models.CASCADE)
+
+profile = Profile()
+reveal_type(profile.user)  # E: Revealed type is 'main.User*'
+[out]
+
+
diff --git a/test/testdjango.py b/test/testdjango.py
index 695562a..d031bb2 100644
--- a/test/testdjango.py
+++ b/test/testdjango.py
@@ -15,7 +15,8 @@ MYPY_INI_PATH = ROOT_DIR / 'test' / 'plugins.ini'
 class DjangoTestSuite(DataSuite):
     files = [
         'check-model-fields.test',
-        'check-postgres-fields.test'
+        'check-postgres-fields.test',
+        'check-model-relations.test'
     ]
     data_prefix = str(TEST_DATA_DIR)