MOTOR-1285 Add generics to type stubs (#272)

doc/contributors.rst

@@ -17,3 +17,4 @@ The following is a list of people who have contributed to
 - Tushar Singh
 - Steven Silvester
 - Julius Park
+- Doeke Buursma

motor/core.pyi

@@ -35,6 +35,7 @@ import pymongo.errors
 import pymongo.mongo_client
 import typing_extensions
 from bson import Binary, Code, CodecOptions, DBRef, Timestamp
+from bson.codec_options import TypeRegistry
 from bson.raw_bson import RawBSONDocument
 from pymongo import IndexModel, ReadPreference, WriteConcern
 from pymongo.change_stream import ChangeStream
@@ -99,7 +100,7 @@ class AgnosticClient(AgnosticBaseProperties[_DocumentType]):
     def __hash__(self) -> int: ...
     async def drop_database(
         self,
-        name_or_database: Union[str, AgnosticDatabase],
+        name_or_database: Union[str, AgnosticDatabase[_DocumentTypeArg]],
         session: Optional[AgnosticClientSession] = None,
         comment: Optional[Any] = None,
     ) -> None: ...
@@ -130,7 +131,7 @@ class AgnosticClient(AgnosticBaseProperties[_DocumentType]):
         session: Optional[AgnosticClientSession] = None,
         comment: Optional[Any] = None,
         **kwargs: Any,
-    ) -> AgnosticCommandCursor: ...
+    ) -> AgnosticCommandCursor[dict[str, Any]]: ...
     async def list_database_names(
         self,
         session: Optional[AgnosticClientSession] = None,
@@ -155,7 +156,16 @@ class AgnosticClient(AgnosticBaseProperties[_DocumentType]):
     _io_loop: Optional[Any]
     _framework: Any
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
+    def __init__(
+        self,
+        host: Optional[Union[str, Sequence[str]]] = None,
+        port: Optional[int] = None,
+        document_class: Optional[type[_DocumentType]] = None,
+        tz_aware: Optional[bool] = None,
+        connect: Optional[bool] = None,
+        type_registry: Optional[TypeRegistry] = None,
+        **kwargs: Any,
+    ) -> None: ...
     @property
     def io_loop(self) -> Any: ...
     def get_io_loop(self) -> Any: ...
@@ -173,9 +183,9 @@ class AgnosticClient(AgnosticBaseProperties[_DocumentType]):
         comment: Optional[str] = None,
         full_document_before_change: Optional[str] = None,
         show_expanded_events: Optional[bool] = None,
-    ) -> AgnosticChangeStream: ...
-    def __getattr__(self, name: str) -> AgnosticDatabase: ...
-    def __getitem__(self, name: str) -> AgnosticDatabase: ...
+    ) -> AgnosticChangeStream[_DocumentType]: ...
+    def __getattr__(self, name: str) -> AgnosticDatabase[_DocumentType]: ...
+    def __getitem__(self, name: str) -> AgnosticDatabase[_DocumentType]: ...
     def wrap(self, obj: Any) -> Any: ...
 
 class _MotorTransactionContext:
@@ -240,7 +250,7 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
         comment: Optional[Any] = None,
         max_await_time_ms: Optional[int] = None,
         **kwargs: Any,
-    ) -> AgnosticCommandCursor: ...
+    ) -> AgnosticCommandCursor[_DocumentType]: ...
     @overload
     async def command(
         self,
@@ -290,7 +300,7 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
         session: Optional[AgnosticClientSession] = None,
         check_exists: Optional[bool] = True,
         **kwargs: Any,
-    ) -> AgnosticCollection: ...
+    ) -> AgnosticCollection[_DocumentType]: ...
     async def dereference(
         self,
         dbref: DBRef,
@@ -300,7 +310,7 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
     ) -> Optional[_DocumentType]: ...
     async def drop_collection(
         self,
-        name_or_collection: Union[str, AgnosticCollection],
+        name_or_collection: Union[str, AgnosticCollection[_DocumentTypeArg]],
         session: Optional[AgnosticClientSession] = None,
         comment: Optional[Any] = None,
         encrypted_fields: Optional[Mapping[str, Any]] = None,
@@ -312,7 +322,7 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> AgnosticCollection: ...
+    ) -> AgnosticCollection[_DocumentType]: ...
     async def list_collection_names(
         self,
         session: Optional[AgnosticClientSession] = None,
@@ -326,12 +336,12 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
         filter: Optional[Mapping[str, Any]] = None,
         comment: Optional[Any] = None,
         **kwargs: Any,
-    ) -> AgnosticCommandCursor: ...
+    ) -> AgnosticCommandCursor[MutableMapping[str, Any]]: ...
     @property
     def name(self) -> str: ...
     async def validate_collection(
         self,
-        name_or_collection: Union[str, AgnosticCollection],
+        name_or_collection: Union[str, AgnosticCollection[_DocumentTypeArg]],
         scandata: bool = False,
         full: bool = False,
         session: Optional[AgnosticClientSession] = None,
@@ -344,14 +354,14 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> AgnosticDatabase: ...
+    ) -> AgnosticDatabase[_DocumentType]: ...
     async def _async_aggregate(
         self, pipeline: _Pipeline, session: Optional[AgnosticClientSession] = None, **kwargs: Any
-    ) -> AgnosticCommandCursor: ...
-    def __init__(self, client: AgnosticClient, name: str, **kwargs: Any) -> None: ...
+    ) -> AgnosticCommandCursor[_DocumentType]: ...
+    def __init__(self, client: AgnosticClient[_DocumentType], name: str, **kwargs: Any) -> None: ...
     def aggregate(
         self, pipeline: _Pipeline, *args: Any, **kwargs: Any
-    ) -> AgnosticLatentCommandCursor: ...
+    ) -> AgnosticLatentCommandCursor[_DocumentType]: ...
     def watch(
         self,
         pipeline: Optional[_Pipeline] = None,
@@ -366,11 +376,11 @@ class AgnosticDatabase(AgnosticBaseProperties[_DocumentType]):
         comment: Optional[Any] = None,
         full_document_before_change: Optional[str] = None,
         show_expanded_events: Optional[bool] = None,
-    ) -> AgnosticChangeStream: ...
+    ) -> AgnosticChangeStream[_DocumentType]: ...
     @property
-    def client(self) -> AgnosticClient: ...
-    def __getattr__(self, name: str) -> AgnosticCollection: ...
-    def __getitem__(self, name: str) -> AgnosticCollection: ...
+    def client(self) -> AgnosticClient[_DocumentType]: ...
+    def __getattr__(self, name: str) -> AgnosticCollection[_DocumentType]: ...
+    def __getitem__(self, name: str) -> AgnosticCollection[_DocumentType]: ...
     def __call__(self, *args: Any, **kwargs: Any) -> None: ...
     def wrap(self, obj: Any) -> Any: ...
     def get_io_loop(self) -> Any: ...
@@ -577,14 +587,14 @@ class AgnosticCollection(AgnosticBaseProperties[_DocumentType]):
         read_preference: Optional[ReadPreference] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional[ReadConcern] = None,
-    ) -> AgnosticCollection[Mapping[str, Any]]: ...
+    ) -> AgnosticCollection[_DocumentType]: ...
     def list_search_indexes(
         self,
         name: Optional[str] = None,
         session: Optional[AgnosticClientSession] = None,
         comment: Optional[Any] = None,
         **kwargs: Any,
-    ) -> AgnosticLatentCommandCursor: ...
+    ) -> AgnosticLatentCommandCursor[Mapping[str, Any]]: ...
     async def create_search_index(
         self,
         model: Union[Mapping[str, SearchIndexModel], Any],
@@ -625,17 +635,19 @@ class AgnosticCollection(AgnosticBaseProperties[_DocumentType]):
         _delegate: Any = None,
         **kwargs: Any,
     ) -> None: ...
-    def __getattr__(self, name: str) -> AgnosticCollection: ...
-    def __getitem__(self, name: str) -> AgnosticCollection: ...
+    def __getattr__(self, name: str) -> AgnosticCollection[_DocumentType]: ...
+    def __getitem__(self, name: str) -> AgnosticCollection[_DocumentType]: ...
     def __call__(self, *args: Any, **kwargs: Any) -> Any: ...
-    def find(self, *args: Any, **kwargs: Any) -> AgnosticCursor: ...
-    def find_raw_batches(self, *args: Any, **kwargs: Any) -> AgnosticCursor: ...
+    def find(self, *args: Any, **kwargs: Any) -> AgnosticCursor[_DocumentType]: ...
+    def find_raw_batches(
+        self, *args: Any, **kwargs: Any
+    ) -> AgnosticRawBatchCursor[_DocumentType]: ...
     def aggregate(
         self, pipeline: _Pipeline, *args: Any, **kwargs: Any
-    ) -> AgnosticCommandCursor: ...
+    ) -> AgnosticCommandCursor[_DocumentType]: ...
     def aggregate_raw_batches(
         self, pipeline: _Pipeline, **kwargs: Any
-    ) -> AgnosticCommandCursor: ...
+    ) -> AgnosticRawBatchCursor[_DocumentType]: ...
     def watch(
         self,
         pipeline: Optional[_Pipeline] = None,
@@ -653,13 +665,17 @@ class AgnosticCollection(AgnosticBaseProperties[_DocumentType]):
     ) -> Any: ...
     def list_indexes(
         self, session: Optional[AgnosticClientSession] = None, **kwargs: Any
-    ) -> AgnosticLatentCommandCursor: ...
+    ) -> AgnosticLatentCommandCursor[MutableMapping[str, Any]]: ...
     def wrap(self, obj: Any) -> Any: ...
     def get_io_loop(self) -> Any: ...
 
-class AgnosticBaseCursor(AgnosticBase):
+class AgnosticBaseCursor(AgnosticBase, Generic[_DocumentType]):
     def __init__(
-        self, cursor: Union[Cursor, CommandCursor, _LatentCursor], collection: AgnosticCollection
+        self,
+        cursor: Union[
+            Cursor[_DocumentType], CommandCursor[_DocumentType], _LatentCursor[_DocumentType]
+        ],
+        collection: AgnosticCollection[_DocumentType],
     ) -> None: ...
     def address(self) -> Optional[_Address]: ...
     def cursor_id(self) -> Optional[int]: ...
@@ -683,72 +699,76 @@ class AgnosticBaseCursor(AgnosticBase):
         self, length: Union[int, None], the_list: list, future: Any, get_more_result: Any
     ) -> None: ...
     def get_io_loop(self) -> Any: ...
-    def batch_size(self, batch_size: int) -> AgnosticBaseCursor: ...
+    def batch_size(self, batch_size: int) -> AgnosticBaseCursor[_DocumentType]: ...
     def _buffer_size(self) -> int: ...
     def _query_flags(self) -> Optional[int]: ...
     def _data(self) -> None: ...
     def _killed(self) -> None: ...
     async def close(self) -> None: ...
 
-class AgnosticCursor(AgnosticBaseCursor):
+class AgnosticCursor(AgnosticBaseCursor[_DocumentType]):
     __motor_class_name__: str
     __delegate_class__: type[Cursor]
-    def collation(self, collation: Optional[_CollationIn]) -> AgnosticCursor: ...
+    def collation(self, collation: Optional[_CollationIn]) -> AgnosticCursor[_DocumentType]: ...
     async def distinct(self, key: str) -> list: ...
     async def explain(self) -> _DocumentType: ...
-    def add_option(self, mask: int) -> AgnosticCursor: ...
-    def remove_option(self, mask: int) -> AgnosticCursor: ...
-    def limit(self, limit: int) -> AgnosticCursor: ...
-    def skip(self, skip: int) -> AgnosticCursor: ...
-    def max_scan(self, max_scan: Optional[int]) -> AgnosticCursor: ...
+    def add_option(self, mask: int) -> AgnosticCursor[_DocumentType]: ...
+    def remove_option(self, mask: int) -> AgnosticCursor[_DocumentType]: ...
+    def limit(self, limit: int) -> AgnosticCursor[_DocumentType]: ...
+    def skip(self, skip: int) -> AgnosticCursor[_DocumentType]: ...
+    def max_scan(self, max_scan: Optional[int]) -> AgnosticCursor[_DocumentType]: ...
     def sort(
         self, key_or_list: _Hint, direction: Optional[Union[int, str]] = None
-    ) -> AgnosticCursor: ...
-    def hint(self, index: Optional[_Hint]) -> AgnosticCursor: ...
-    def where(self, code: Union[str, Code]) -> AgnosticCursor: ...
-    def max_await_time_ms(self, max_await_time_ms: Optional[int]) -> AgnosticCursor: ...
-    def max_time_ms(self, max_time_ms: Optional[int]) -> AgnosticCursor: ...
-    def min(self, spec: _Sort) -> AgnosticCursor: ...
-    def max(self, spec: _Sort) -> AgnosticCursor: ...
-    def comment(self, comment: Any) -> AgnosticCursor: ...
-    def allow_disk_use(self, allow_disk_use: bool) -> AgnosticCursor: ...
-    def rewind(self) -> AgnosticCursor: ...
-    def clone(self) -> AgnosticCursor: ...
-    def __copy__(self) -> AgnosticCursor: ...
-    def __deepcopy__(self, memo: Any) -> AgnosticCursor: ...
+    ) -> AgnosticCursor[_DocumentType]: ...
+    def hint(self, index: Optional[_Hint]) -> AgnosticCursor[_DocumentType]: ...
+    def where(self, code: Union[str, Code]) -> AgnosticCursor[_DocumentType]: ...
+    def max_await_time_ms(
+        self, max_await_time_ms: Optional[int]
+    ) -> AgnosticCursor[_DocumentType]: ...
+    def max_time_ms(self, max_time_ms: Optional[int]) -> AgnosticCursor[_DocumentType]: ...
+    def min(self, spec: _Sort) -> AgnosticCursor[_DocumentType]: ...
+    def max(self, spec: _Sort) -> AgnosticCursor[_DocumentType]: ...
+    def comment(self, comment: Any) -> AgnosticCursor[_DocumentType]: ...
+    def allow_disk_use(self, allow_disk_use: bool) -> AgnosticCursor[_DocumentType]: ...
+    def rewind(self) -> AgnosticCursor[_DocumentType]: ...
+    def clone(self) -> AgnosticCursor[_DocumentType]: ...
+    def __copy__(self) -> AgnosticCursor[_DocumentType]: ...
+    def __deepcopy__(self, memo: Any) -> AgnosticCursor[_DocumentType]: ...
     def _query_flags(self) -> int: ...
     def _data(self) -> Any: ...
     def _killed(self) -> Any: ...
 
-class AgnosticRawBatchCursor(AgnosticCursor):
+class AgnosticRawBatchCursor(AgnosticCursor[_DocumentType]):
     __motor_class_name__: str
     __delegate_class__: type[RawBatchCursor]
 
-class AgnosticCommandCursor(AgnosticBaseCursor):
+class AgnosticCommandCursor(AgnosticBaseCursor[_DocumentType]):
     __motor_class_name__: str
     __delegate_class__: type[CommandCursor]
 
     def _query_flags(self) -> int: ...
     def _data(self) -> Any: ...
     def _killed(self) -> Any: ...
 
-class AgnosticRawBatchCommandCursor(AgnosticCommandCursor):
+class AgnosticRawBatchCommandCursor(AgnosticCommandCursor[_DocumentType]):
     __motor_class_name__: str
     __delegate_class__: type[RawBatchCommandCursor]
 
-class _LatentCursor:
-    def __init__(self, collection: AgnosticCollection): ...
+class _LatentCursor(Generic[_DocumentType]):
+    def __init__(self, collection: AgnosticCollection[_DocumentType]): ...
     def _CommandCursor__end_session(self, *args: Any, **kwargs: Any) -> None: ...
     def _CommandCursor__die(self, *args: Any, **kwargs: Any) -> None: ...
-    def clone(self) -> _LatentCursor: ...
-    def rewind(self) -> _LatentCursor: ...
+    def clone(self) -> _LatentCursor[_DocumentType]: ...
+    def rewind(self) -> _LatentCursor[_DocumentType]: ...
 
-class AgnosticLatentCommandCursor(AgnosticCommandCursor):
+class AgnosticLatentCommandCursor(AgnosticCommandCursor[_DocumentType]):
     __motor_class_name__: str
-    def __init__(self, collection: AgnosticCollection, start: Any, *args: Any, **kwargs: Any): ...
+    def __init__(
+        self, collection: AgnosticCollection[_DocumentType], start: Any, *args: Any, **kwargs: Any
+    ): ...
     def _on_started(self, original_future: Any, future: Any) -> None: ...
 
-class AgnosticChangeStream(AgnosticBase):
+class AgnosticChangeStream(AgnosticBase, Generic[_DocumentType]):
     __motor_class_name__: str
     __delegate_class__: type[ChangeStream]
 
@@ -778,22 +798,22 @@ class AgnosticChangeStream(AgnosticBase):
     async def next(self) -> _DocumentType: ...
     async def try_next(self) -> Optional[_DocumentType]: ...
     async def close(self) -> None: ...
-    def __aiter__(self) -> AgnosticChangeStream: ...
+    def __aiter__(self) -> AgnosticChangeStream[_DocumentType]: ...
     __anext__ = next
-    async def __aenter__(self) -> AgnosticChangeStream: ...
+    async def __aenter__(self) -> AgnosticChangeStream[_DocumentType]: ...
     async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: ...
     def get_io_loop(self) -> Any: ...
     def __enter__(self) -> None: ...
     def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: ...
 
-class AgnosticClientEncryption(AgnosticBase):
+class AgnosticClientEncryption(AgnosticBase, Generic[_DocumentType]):
     __motor_class_name__: str
     __delegate_class__: type[ClientEncryption]
     def __init__(
         self,
         kms_providers: Mapping[str, Any],
         key_vault_namespace: str,
-        key_vault_client: AgnosticClient,
+        key_vault_client: AgnosticClient[_DocumentTypeArg],
         codec_options: CodecOptions,
         io_loop: Optional[Any] = None,
         kms_tls_options: Optional[Mapping[str, Any]] = None,
@@ -843,17 +863,17 @@ class AgnosticClientEncryption(AgnosticBase):
     @property
     def io_loop(self) -> Any: ...
     def get_io_loop(self) -> Any: ...
-    async def __aenter__(self) -> AgnosticClientEncryption: ...
+    async def __aenter__(self) -> AgnosticClientEncryption[_DocumentType]: ...
     async def __aexit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: ...
     def __enter__(self) -> NoReturn: ...
     def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: ...
-    async def get_keys(self) -> AgnosticCursor: ...
+    async def get_keys(self) -> AgnosticCursor[RawBSONDocument]: ...
     async def create_encrypted_collection(
         self,
-        database: AgnosticDatabase,
+        database: AgnosticDatabase[_DocumentTypeArg],
         name: str,
         encrypted_fields: Mapping[str, Any],
         kms_provider: Optional[str] = None,
         master_key: Optional[Mapping[str, Any]] = None,
         **kwargs: Any,
-    ) -> tuple[AgnosticCollection, Mapping[str, Any]]: ...
+    ) -> tuple[AgnosticCollection[_DocumentTypeArg], Mapping[str, Any]]: ...