From c8c5f8b51b5cd4d46af01a52357c34aa89c812f9 Mon Sep 17 00:00:00 2001
From: Maxim Kurnikov <maxim.kurnikov@gmail.com>
Date: Tue, 2 Oct 2018 02:20:44 +0300
Subject: [PATCH] stubs for manager/testcases

---
 .gitignore                         |   3 +-
 django-stubs/db/models/manager.pyi | 115 +++++++++++++++++++++++++++--
 django-stubs/test/__init__.pyi     |   5 ++
 3 files changed, 117 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index 6b79feb..25e61e2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,4 +2,5 @@
 __pycache__/
 out/
 /test_sqlite.py
-/django
\ No newline at end of file
+/django
+.idea/
\ No newline at end of file
diff --git a/django-stubs/db/models/manager.pyi b/django-stubs/db/models/manager.pyi
index 526bcd6..9809424 100644
--- a/django-stubs/db/models/manager.pyi
+++ b/django-stubs/db/models/manager.pyi
@@ -1,8 +1,15 @@
-from typing import Any, Dict, List, Optional, Tuple, Type
+from collections import OrderedDict
+from datetime import date, datetime
+from decimal import Decimal
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, TypeVar, Set, Generic
+from unittest.mock import MagicMock
 
 from django.contrib.sites.managers import CurrentSiteManager
+from django.db.models import Q
 from django.db.models.base import Model
-from django.db.models.query import QuerySet
+from django.db.models.query import QuerySet, RawQuerySet
+
+_T = TypeVar('_T', bound=Model)
 
 
 class BaseManager:
@@ -32,10 +39,108 @@ class BaseManager:
     def __eq__(self, other: Optional[CurrentSiteManager]) -> bool: ...
     def __hash__(self): ...
 
-class Manager: ...
+class Manager(Generic[_T]):
+    def exists(self) -> bool: ...
+
+    def explain(
+            self, *, format: Optional[Any] = ..., **options: Any
+    ) -> str: ...
+
+    def raw(
+            self,
+            raw_query: str,
+            params: Optional[
+                Union[
+                    Dict[str, str],
+                    List[datetime],
+                    List[Decimal],
+                    List[str],
+                    Set[str],
+                    Tuple[int],
+                ]
+            ] = ...,
+            translations: Optional[Dict[str, str]] = ...,
+            using: None = ...,
+    ) -> RawQuerySet: ...
+
+    def values(self, *fields: Any, **expressions: Any) -> QuerySet: ...
+
+    def values_list(
+            self, *fields: Any, flat: bool = ..., named: bool = ...
+    ) -> QuerySet: ...
+
+    def dates(
+            self, field_name: str, kind: str, order: str = ...
+    ) -> QuerySet: ...
+
+    def datetimes(
+            self, field_name: str, kind: str, order: str = ..., tzinfo: None = ...
+    ) -> QuerySet: ...
+
+    def none(self) -> QuerySet[_T]: ...
+
+    def all(self) -> QuerySet[_T]: ...
+
+    def filter(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
+
+    def exclude(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
+
+    def complex_filter(
+            self,
+            filter_obj: Union[
+                Dict[str, datetime], Dict[str, QuerySet], Q, MagicMock
+            ],
+    ) -> QuerySet[_T]: ...
+
+    def union(self, *other_qs: Any, all: bool = ...) -> QuerySet[_T]: ...
+
+    def intersection(self, *other_qs: Any) -> QuerySet[_T]: ...
+
+    def difference(self, *other_qs: Any) -> QuerySet[_T]: ...
+
+    def select_for_update(
+            self, nowait: bool = ..., skip_locked: bool = ..., of: Tuple = ...
+    ) -> QuerySet: ...
+
+    def select_related(self, *fields: Any) -> QuerySet[_T]: ...
+
+    def prefetch_related(self, *lookups: Any) -> QuerySet[_T]: ...
+
+    def annotate(self, *args: Any, **kwargs: Any) -> QuerySet[_T]: ...
+
+    def order_by(self, *field_names: Any) -> QuerySet[_T]: ...
+
+    def distinct(self, *field_names: Any) -> QuerySet[_T]: ...
+
+    def extra(
+            self,
+            select: Optional[
+                Union[Dict[str, int], Dict[str, str], OrderedDict]
+            ] = ...,
+            where: Optional[List[str]] = ...,
+            params: Optional[Union[List[int], List[str]]] = ...,
+            tables: Optional[List[str]] = ...,
+            order_by: Optional[Union[List[str], Tuple[str]]] = ...,
+            select_params: Optional[Union[List[int], List[str], Tuple[int]]] = ...,
+    ) -> QuerySet[_T]: ...
+
+    def iterator(self, chunk_size: int = ...) -> Iterator[_T]: ...
+
+    def aggregate(
+            self, *args: Any, **kwargs: Any
+    ) -> Dict[str, Optional[Union[datetime, float]]]: ...
+
+    def count(self) -> int: ...
+
+    def get(
+            self, *args: Any, **kwargs: Any
+    ) -> _T: ...
+
+    def create(self, **kwargs: Any) -> _T: ...
+
 
 class ManagerDescriptor:
-    manager: django.db.models.manager.Manager = ...
+    manager: Manager = ...
     def __init__(self, manager: Manager) -> None: ...
     def __get__(
         self, instance: Optional[Model], cls: Type[Model] = ...
@@ -44,6 +149,6 @@ class ManagerDescriptor:
 class EmptyManager(Manager):
     creation_counter: int
     name: None
-    model: Optional[Type[django.db.models.base.Model]] = ...
+    model: Optional[Type[Model]] = ...
     def __init__(self, model: Type[Model]) -> None: ...
     def get_queryset(self) -> QuerySet: ...
diff --git a/django-stubs/test/__init__.pyi b/django-stubs/test/__init__.pyi
index e69de29..14d6026 100644
--- a/django-stubs/test/__init__.pyi
+++ b/django-stubs/test/__init__.pyi
@@ -0,0 +1,5 @@
+from .testcases import (
+    TestCase as TestCase,
+    TransactionTestCase as TransactionTestCase,
+    SimpleTestCase as SimpleTestCase
+)
\ No newline at end of file