add(mongoengine): missing queryset methods (#8)

Some notable ones: modify, upsert_one, explain

Fix some param names and types.

Also add flake8-pyi to check the stubs for errors, but pyright might already cover it.

pyproject.toml

@@ -25,6 +25,7 @@ flake8 = "^3.8"
 black = {version = "20.8b1",allows-prereleases = true}
 isort = "^4.3"
 wheel = "^0.36.2"
+flake8-pyi = "^20.10"
 
 [build-system]
 requires = ["poetry>=0.12", "setuptools"]

tests/test_mongoengine.py

@@ -118,12 +118,18 @@ def main() -> None:
     # reveal_type(posts)
     # Revealed type is 'builtins.list[test_mongoengine.Post*]
 
+    posts_again = Post.objects().exclude("title").all_fields()
+    print(posts_again)
+
     insert_result = Post.objects().insert([Post(), Post()])
     print(insert_result)
 
     insert_result_2 = Post.objects().insert([Post(), Post()], load_bulk=False)
     print(insert_result_2)
 
+    res: int = Post.objects().update(foo=True)
+    print(res)
+
     first_post = posts[0]
     first_post.tags.values()
     first_post.errors

typings/mongoengine/queryset/queryset.pyi

@@ -15,15 +15,19 @@ from typing import (
 
 from bson import ObjectId
 from mongoengine.queryset.visitor import Q
+from pymongo.collation import Collation
 from pymongo.collection import Collection
 from pymongo.command_cursor import CommandCursor
 from pymongo.cursor import Cursor
+from pymongo.read_preferences import _ServerMode
 from typing_extensions import Literal
 
 _T = TypeVar("_T")
 
 _U = TypeVar("_U", bound="QuerySet[Any]")
 
+_ReadWriteConcern = Dict[str, Union[str, int, bool]]
+
 class QuerySet(Generic[_T]):
     _document: Type[_T]
     _collection: Collection
@@ -40,38 +44,76 @@ class QuerySet(Generic[_T]):
         self,
         doc_or_docs: Union[_T, List[_T]],
         load_bulk: Literal[False],
-        write_concern: Optional[Any] = ...,
+        write_concern: Optional[_ReadWriteConcern] = ...,
         signal_kwargs: Optional[Any] = ...,
     ) -> List[ObjectId]: ...
     @overload
     def insert(
         self,
         doc_or_docs: Union[_T, List[_T]],
         load_bulk: Literal[True] = ...,
-        write_concern: Optional[Any] = ...,
+        write_concern: Optional[_ReadWriteConcern] = ...,
         signal_kwargs: Optional[Any] = ...,
     ) -> List[_T]: ...
     def values(self, *args: str) -> Dict[str, Any]: ...
     def as_pymongo(self) -> QuerySet[Dict[str, Any]]: ...
+    def scalar(self, *fields: str) -> List[Any]: ...
     def values_list(self, *args: str) -> List[Any]: ...
-    def update(self, *args: Q, **kwargs: object) -> int: ...
+    def update(
+        self,
+        upsert: bool = ...,
+        multi: bool = ...,
+        write_concern: Optional[_ReadWriteConcern] = ...,
+        read_concern: Optional[_ReadWriteConcern] = ...,
+        **update: object
+    ) -> int: ...
     def update_one(
-        self, upsert: bool = ..., write_concern: Optional[Any] = ..., **update: object
+        self,
+        upsert: bool = ...,
+        write_concern: Optional[_ReadWriteConcern] = ...,
+        **update: object
     ) -> int: ...
+    def upsert_one(
+        self,
+        write_concern: Optional[_ReadWriteConcern] = ...,
+        read_concern: Optional[_ReadWriteConcern] = ...,
+        **update: object
+    ) -> _T: ...
     def delete(self) -> None: ...
     def filter(self: _U, *args: Q, **kwargs: object) -> _U: ...
     def none(self: _U) -> _U: ...
-    def read_preference(self: _U, read_preference: object) -> _U: ...
-    def only(self: _U, *args: str) -> _U: ...
-    def order_by(self: _U, *args: str) -> _U: ...
+    def read_preference(self: _U, read_preference: _ServerMode) -> _U: ...
+    def rewind(self: _U) -> _U: ...
+    def only(self: _U, *fields: str) -> _U: ...
+    def exclude(self: _U, *args: str) -> _U: ...
+    def exec_js(self: _U, code: str, *fields: str, **options: Any) -> _U: ...
+    def explain(self) -> Any: ...
+    def all_fields(self: _U) -> _U: ...
+    def average(self: _U, field: str) -> _U: ...
+    def batch_size(self: _U, size: int) -> _U: ...
+    def clone(self: _U) -> _U: ...
+    def collation(self: _U, collation: Optional[Collation] = ...) -> _U: ...
+    def comment(self: _U, text: str) -> _U: ...
+    def order_by(self: _U, *keys: str) -> _U: ...
     def skip(self: _U, amount: Optional[int]) -> _U: ...
+    def sum(self, field: str) -> float: ...
     def limit(self: _U, amount: int) -> _U: ...
     def count(self) -> int: ...
     def __iter__(self) -> Iterator[_T]: ...
     def __next__(self) -> _T: ...
     def all(self: _U) -> _U: ...
     def aggregate(self, *args: object, allowDiskUse: bool = ...) -> CommandCursor: ...
     def max_time_ms(self: _U, ms: Optional[int]) -> _U: ...
+    def modify(
+        self: _U,
+        upsert: bool = ...,
+        full_response: bool = ...,
+        remove: bool = ...,
+        new: bool = ...,
+        **update: object
+    ) -> _U: ...
+    def no_cache(self: _U) -> _U: ...
+    def no_dereference(self: _U) -> _U: ...
     def create(self, **kwargs: object) -> _T: ...
     def distinct(self, field: str) -> List[Any]: ...
     def __getitem__(self, key: int) -> _T: ...

typings/pymongo/read_preferences.pyi

@@ -1,6 +1,13 @@
+class _ServerMode: ...
+class Primary(_ServerMode): ...
+class PrimaryPreferred(_ServerMode): ...
+class Secondary(_ServerMode): ...
+class SecondaryPreferred(_ServerMode): ...
+class Nearest(_ServerMode): ...
+
 class ReadPreference:
-    NEAREST: object
-    PRIMARY: object
-    PRIMARY_PREFERRED: object
-    SECONDARY: object
-    SECONDARY_PREFERRED: object
+    NEAREST: Nearest
+    PRIMARY: Primary
+    PRIMARY_PREFERRED: PrimaryPreferred
+    SECONDARY: Secondary
+    SECONDARY_PREFERRED: SecondaryPreferred