From 4ba7055238e9ddf1f44f0c3993cf0283351f99da Mon Sep 17 00:00:00 2001
From: collerek <collerek@gmail.com>
Date: Mon, 22 Mar 2021 16:16:26 +0100
Subject: [PATCH] add relation fields overloads

---
 ormar/fields/foreign_key.py  | 28 ++++++++++++++++++++++++----
 ormar/fields/many_to_many.py | 28 ++++++++++++++++++++++++----
 2 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/ormar/fields/foreign_key.py b/ormar/fields/foreign_key.py
index 65a5311..0c1894d 100644
--- a/ormar/fields/foreign_key.py
+++ b/ormar/fields/foreign_key.py
@@ -3,7 +3,17 @@ import sys
 import uuid
 from dataclasses import dataclass
 from random import choices
-from typing import Any, Dict, List, Optional, TYPE_CHECKING, Tuple, Type, Union
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    TYPE_CHECKING,
+    Tuple,
+    Type,
+    Union,
+    overload,
+)
 
 import sqlalchemy
 from pydantic import BaseModel, create_model
@@ -19,9 +29,9 @@ if TYPE_CHECKING:  # pragma no cover
     from ormar.fields import ManyToManyField
 
     if sys.version_info < (3, 7):
-        ToType = Type["Model"]
+        ToType = Type["T"]
     else:
-        ToType = Union[Type["Model"], "ForwardRef"]
+        ToType = Union[Type["T"], "ForwardRef"]
 
 
 def create_dummy_instance(fk: Type["T"], pk: Any = None) -> "T":
@@ -168,8 +178,18 @@ class ForeignKeyConstraint:
     onupdate: Optional[str]
 
 
+@overload
+def ForeignKey(to: Type["T"], **kwargs: Any) -> "T":
+    ...
+
+
+@overload
+def ForeignKey(to: ForwardRef, **kwargs: Any) -> "Model":
+    ...
+
+
 def ForeignKey(  # noqa CFQ002
-    to: Type["T"],
+    to: "ToType",
     *,
     name: str = None,
     unique: bool = False,
diff --git a/ormar/fields/many_to_many.py b/ormar/fields/many_to_many.py
index a0989e0..d488029 100644
--- a/ormar/fields/many_to_many.py
+++ b/ormar/fields/many_to_many.py
@@ -1,5 +1,15 @@
 import sys
-from typing import Any, List, Optional, TYPE_CHECKING, Tuple, Type, Union, cast
+from typing import (
+    Any,
+    List,
+    Optional,
+    TYPE_CHECKING,
+    Tuple,
+    Type,
+    Union,
+    cast,
+    overload,
+)
 
 from pydantic.typing import ForwardRef, evaluate_forwardref
 import ormar  # noqa: I100
@@ -12,9 +22,9 @@ if TYPE_CHECKING:  # pragma no cover
     from ormar.relations.relation_proxy import RelationProxy
 
     if sys.version_info < (3, 7):
-        ToType = Type["Model"]
+        ToType = Type["T"]
     else:
-        ToType = Union[Type["Model"], "ForwardRef"]
+        ToType = Union[Type["T"], "ForwardRef"]
 
 REF_PREFIX = "#/components/schemas/"
 
@@ -58,8 +68,18 @@ def populate_m2m_params_based_on_to_model(
     return __type__, column_type
 
 
+@overload
+def ManyToMany(to: Type["T"], **kwargs: Any) -> "RelationProxy[T]":
+    ...
+
+
+@overload
+def ManyToMany(to: ForwardRef, **kwargs: Any) -> "RelationProxy":
+    ...
+
+
 def ManyToMany(
-    to: Type["T"],
+    to: "ToType",
     through: Optional["ToType"] = None,
     *,
     name: str = None,