From 9175f6004cf1eeed4b3feec8ed21b90727aaa211 Mon Sep 17 00:00:00 2001
From: collerek <collerek@gmail.com>
Date: Sun, 1 Nov 2020 08:29:06 +0100
Subject: [PATCH] refactor many2many typehints into protocols

---
 ormar/__init__.py                    |  8 ++--
 ormar/fields/many_to_many.py         | 61 ++--------------------------
 ormar/protocols/__init__.py          |  4 ++
 ormar/protocols/queryset_protocol.py | 41 +++++++++++++++++++
 ormar/protocols/relation_protocol.py | 12 ++++++
 ormar/relations/querysetproxy.py     |  2 +-
 6 files changed, 66 insertions(+), 62 deletions(-)
 create mode 100644 ormar/protocols/__init__.py
 create mode 100644 ormar/protocols/queryset_protocol.py
 create mode 100644 ormar/protocols/relation_protocol.py

diff --git a/ormar/__init__.py b/ormar/__init__.py
index ecead7c..17324ad 100644
--- a/ormar/__init__.py
+++ b/ormar/__init__.py
@@ -1,5 +1,6 @@
 from ormar.exceptions import ModelDefinitionError, ModelNotSet, MultipleMatches, NoMatch
-from ormar.fields import (
+from ormar.protocols import QuerySetProtocol, RelationProtocol  # noqa: I100
+from ormar.fields import (  # noqa: I100
     BigInteger,
     Boolean,
     Date,
@@ -28,8 +29,7 @@ class UndefinedType:  # pragma no cover
 
 Undefined = UndefinedType()
 
-
-__version__ = "0.3.11"
+__version__ = "0.4.0"
 __all__ = [
     "Integer",
     "BigInteger",
@@ -54,4 +54,6 @@ __all__ = [
     "Undefined",
     "UUID",
     "UniqueColumns",
+    "QuerySetProtocol",
+    "RelationProtocol",
 ]
diff --git a/ormar/fields/many_to_many.py b/ormar/fields/many_to_many.py
index 79fb7fa..0573a86 100644
--- a/ormar/fields/many_to_many.py
+++ b/ormar/fields/many_to_many.py
@@ -1,5 +1,6 @@
-from typing import Any, List, Optional, Sequence, TYPE_CHECKING, Type, Union
+from typing import Any, List, Optional, TYPE_CHECKING, Type, Union
 
+import ormar
 from ormar.fields import BaseField
 from ormar.fields.foreign_key import ForeignKeyField
 
@@ -47,61 +48,5 @@ def ManyToMany(
     return type("ManyToMany", (ManyToManyField, BaseField), namespace)
 
 
-class ManyToManyField(ForeignKeyField):
+class ManyToManyField(ForeignKeyField, ormar.QuerySetProtocol, ormar.RelationProtocol):
     through: Type["Model"]
-
-    if TYPE_CHECKING:  # noqa: C901; #pragma nocover
-
-        @staticmethod
-        async def add(item: "Model") -> None:
-            pass
-
-        @staticmethod
-        async def remove(item: "Model") -> None:
-            pass
-
-        from ormar import QuerySet
-
-        @staticmethod
-        def filter(**kwargs: Any) -> "QuerySet":  # noqa: A003, A001
-            pass
-
-        @staticmethod
-        def select_related(related: Union[List, str]) -> "QuerySet":
-            pass
-
-        @staticmethod
-        async def exists() -> bool:
-            pass
-
-        @staticmethod
-        async def count() -> int:
-            pass
-
-        @staticmethod
-        async def clear() -> int:
-            pass
-
-        @staticmethod
-        def limit(limit_count: int) -> "QuerySet":
-            pass
-
-        @staticmethod
-        def offset(offset: int) -> "QuerySet":
-            pass
-
-        @staticmethod
-        async def first(**kwargs: Any) -> "Model":
-            pass
-
-        @staticmethod
-        async def get(**kwargs: Any) -> "Model":
-            pass
-
-        @staticmethod
-        async def all(**kwargs: Any) -> Sequence[Optional["Model"]]:  # noqa: A003, A001
-            pass
-
-        @staticmethod
-        async def create(**kwargs: Any) -> "Model":
-            pass
diff --git a/ormar/protocols/__init__.py b/ormar/protocols/__init__.py
new file mode 100644
index 0000000..214cc10
--- /dev/null
+++ b/ormar/protocols/__init__.py
@@ -0,0 +1,4 @@
+from ormar.protocols.queryset_protocol import QuerySetProtocol
+from ormar.protocols.relation_protocol import RelationProtocol
+
+__all__ = ["QuerySetProtocol", "RelationProtocol"]
diff --git a/ormar/protocols/queryset_protocol.py b/ormar/protocols/queryset_protocol.py
new file mode 100644
index 0000000..e1c1284
--- /dev/null
+++ b/ormar/protocols/queryset_protocol.py
@@ -0,0 +1,41 @@
+from typing import Any, List, Optional, Protocol, Sequence, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:  # noqa: C901; #pragma nocover
+    from ormar import QuerySet, Model
+
+
+class QuerySetProtocol(Protocol):  # pragma: nocover
+    def filter(self, **kwargs: Any) -> "QuerySet":  # noqa: A003, A001
+        ...
+
+    def select_related(self, related: Union[List, str]) -> "QuerySet":
+        ...
+
+    async def exists(self) -> bool:
+        ...
+
+    async def count(self) -> int:
+        ...
+
+    async def clear(self) -> int:
+        ...
+
+    def limit(self, limit_count: int) -> "QuerySet":
+        ...
+
+    def offset(self, offset: int) -> "QuerySet":
+        ...
+
+    async def first(self, **kwargs: Any) -> "Model":
+        ...
+
+    async def get(self, **kwargs: Any) -> "Model":
+        ...
+
+    async def all(  # noqa: A003, A001
+        self, **kwargs: Any
+    ) -> Sequence[Optional["Model"]]:
+        ...
+
+    async def create(self, **kwargs: Any) -> "Model":
+        ...
diff --git a/ormar/protocols/relation_protocol.py b/ormar/protocols/relation_protocol.py
new file mode 100644
index 0000000..ec0fe98
--- /dev/null
+++ b/ormar/protocols/relation_protocol.py
@@ -0,0 +1,12 @@
+from typing import Protocol, TYPE_CHECKING, Type, Union
+
+if TYPE_CHECKING:  # pragma: nocover
+    from ormar import Model
+
+
+class RelationProtocol(Protocol):  # pragma: nocover
+    def add(self, child: "Model") -> None:
+        ...
+
+    def remove(self, child: Union["Model", Type["Model"]]) -> None:
+        ...
diff --git a/ormar/relations/querysetproxy.py b/ormar/relations/querysetproxy.py
index 88c8dc1..8e2d6b4 100644
--- a/ormar/relations/querysetproxy.py
+++ b/ormar/relations/querysetproxy.py
@@ -10,7 +10,7 @@ if TYPE_CHECKING:  # pragma no cover
     T = TypeVar("T", bound=Model)
 
 
-class QuerysetProxy:
+class QuerysetProxy(ormar.QuerySetProtocol):
     if TYPE_CHECKING:  # pragma no cover
         relation: "Relation"