WIP

Author: NoahStapp

gridfs/asynchronous/grid_file.py

@@ -70,7 +70,7 @@ def _disallow_transactions(session: Optional[AsyncClientSession]) -> None:
 class AsyncGridFS:
     """An instance of GridFS on top of a single Database."""
 
-    def __init__(self, database: AsyncDatabase, collection: str = "fs"):
+    def __init__(self, database: AsyncDatabase[Any], collection: str = "fs"):
         """Create a new instance of :class:`GridFS`.
 
         Raises :class:`TypeError` if `database` is not an instance of
@@ -463,7 +463,7 @@ class AsyncGridFSBucket:
 
     def __init__(
         self,
-        db: AsyncDatabase,
+        db: AsyncDatabase[Any],
         bucket_name: str = "fs",
         chunk_size_bytes: int = DEFAULT_CHUNK_SIZE,
         write_concern: Optional[WriteConcern] = None,
@@ -513,11 +513,11 @@ def __init__(
 
         self._bucket_name = bucket_name
         self._collection = db[bucket_name]
-        self._chunks: AsyncCollection = self._collection.chunks.with_options(
+        self._chunks: AsyncCollection[Any] = self._collection.chunks.with_options(
             write_concern=write_concern, read_preference=read_preference
         )
 
-        self._files: AsyncCollection = self._collection.files.with_options(
+        self._files: AsyncCollection[Any] = self._collection.files.with_options(
             write_concern=write_concern, read_preference=read_preference
         )
 
@@ -1085,7 +1085,7 @@ class AsyncGridIn:
 
     def __init__(
         self,
-        root_collection: AsyncCollection,
+        root_collection: AsyncCollection[Any],
         session: Optional[AsyncClientSession] = None,
         **kwargs: Any,
     ) -> None:
@@ -1141,7 +1141,7 @@ def __init__(
         """
         if not isinstance(root_collection, AsyncCollection):
             raise TypeError(
-                f"root_collection must be an instance of AsyncCollection, not {type(root_collection)}"
+                f"root_collection must be an instance of AsyncCollection[Any], not {type(root_collection)}"
             )
 
         if not root_collection.write_concern.acknowledged:
@@ -1172,7 +1172,7 @@ def __init__(
         object.__setattr__(self, "_buffered_docs_size", 0)
 
     async def _create_index(
-        self, collection: AsyncCollection, index_key: Any, unique: bool
+        self, collection: AsyncCollection[Any], index_key: Any, unique: bool
     ) -> None:
         doc = await collection.find_one(projection={"_id": 1}, session=self._session)
         if doc is None:
@@ -1456,7 +1456,7 @@ class AsyncGridOut(GRIDOUT_BASE_CLASS):  # type: ignore
 
     def __init__(
         self,
-        root_collection: AsyncCollection,
+        root_collection: AsyncCollection[Any],
         file_id: Optional[int] = None,
         file_document: Optional[Any] = None,
         session: Optional[AsyncClientSession] = None,
@@ -1494,7 +1494,7 @@ def __init__(
         """
         if not isinstance(root_collection, AsyncCollection):
             raise TypeError(
-                f"root_collection must be an instance of AsyncCollection, not {type(root_collection)}"
+                f"root_collection must be an instance of AsyncCollection[Any], not {type(root_collection)}"
             )
         _disallow_transactions(session)
 
@@ -1829,7 +1829,7 @@ class _AsyncGridOutChunkIterator:
     def __init__(
         self,
         grid_out: AsyncGridOut,
-        chunks: AsyncCollection,
+        chunks: AsyncCollection[Any],
         session: Optional[AsyncClientSession],
         next_chunk: Any,
     ) -> None:
@@ -1842,7 +1842,7 @@ def __init__(
         self._num_chunks = math.ceil(float(self._length) / self._chunk_size)
         self._cursor = None
 
-    _cursor: Optional[AsyncCursor]
+    _cursor: Optional[AsyncCursor[Any]]
 
     def expected_chunk_length(self, chunk_n: int) -> int:
         if chunk_n < self._num_chunks - 1:
@@ -1921,7 +1921,7 @@ async def close(self) -> None:
 
 class AsyncGridOutIterator:
     def __init__(
-        self, grid_out: AsyncGridOut, chunks: AsyncCollection, session: AsyncClientSession
+        self, grid_out: AsyncGridOut, chunks: AsyncCollection[Any], session: AsyncClientSession
     ):
         self._chunk_iter = _AsyncGridOutChunkIterator(grid_out, chunks, session, 0)
 
@@ -1935,14 +1935,14 @@ async def next(self) -> bytes:
     __anext__ = next
 
 
-class AsyncGridOutCursor(AsyncCursor):
+class AsyncGridOutCursor(AsyncCursor):  # type: ignore[type-arg]
     """A cursor / iterator for returning GridOut objects as the result
     of an arbitrary query against the GridFS files collection.
     """
 
     def __init__(
         self,
-        collection: AsyncCollection,
+        collection: AsyncCollection[Any],
         filter: Optional[Mapping[str, Any]] = None,
         skip: int = 0,
         limit: int = 0,

gridfs/synchronous/grid_file.py

@@ -70,7 +70,7 @@ def _disallow_transactions(session: Optional[ClientSession]) -> None:
 class GridFS:
     """An instance of GridFS on top of a single Database."""
 
-    def __init__(self, database: Database, collection: str = "fs"):
+    def __init__(self, database: Database[Any], collection: str = "fs"):
         """Create a new instance of :class:`GridFS`.
 
         Raises :class:`TypeError` if `database` is not an instance of
@@ -461,7 +461,7 @@ class GridFSBucket:
 
     def __init__(
         self,
-        db: Database,
+        db: Database[Any],
         bucket_name: str = "fs",
         chunk_size_bytes: int = DEFAULT_CHUNK_SIZE,
         write_concern: Optional[WriteConcern] = None,
@@ -511,11 +511,11 @@ def __init__(
 
         self._bucket_name = bucket_name
         self._collection = db[bucket_name]
-        self._chunks: Collection = self._collection.chunks.with_options(
+        self._chunks: Collection[Any] = self._collection.chunks.with_options(
             write_concern=write_concern, read_preference=read_preference
         )
 
-        self._files: Collection = self._collection.files.with_options(
+        self._files: Collection[Any] = self._collection.files.with_options(
             write_concern=write_concern, read_preference=read_preference
         )
 
@@ -1077,7 +1077,7 @@ class GridIn:
 
     def __init__(
         self,
-        root_collection: Collection,
+        root_collection: Collection[Any],
         session: Optional[ClientSession] = None,
         **kwargs: Any,
     ) -> None:
@@ -1133,7 +1133,7 @@ def __init__(
         """
         if not isinstance(root_collection, Collection):
             raise TypeError(
-                f"root_collection must be an instance of Collection, not {type(root_collection)}"
+                f"root_collection must be an instance of Collection[Any], not {type(root_collection)}"
             )
 
         if not root_collection.write_concern.acknowledged:
@@ -1163,7 +1163,7 @@ def __init__(
         object.__setattr__(self, "_buffered_docs", [])
         object.__setattr__(self, "_buffered_docs_size", 0)
 
-    def _create_index(self, collection: Collection, index_key: Any, unique: bool) -> None:
+    def _create_index(self, collection: Collection[Any], index_key: Any, unique: bool) -> None:
         doc = collection.find_one(projection={"_id": 1}, session=self._session)
         if doc is None:
             try:
@@ -1444,7 +1444,7 @@ class GridOut(GRIDOUT_BASE_CLASS):  # type: ignore
 
     def __init__(
         self,
-        root_collection: Collection,
+        root_collection: Collection[Any],
         file_id: Optional[int] = None,
         file_document: Optional[Any] = None,
         session: Optional[ClientSession] = None,
@@ -1482,7 +1482,7 @@ def __init__(
         """
         if not isinstance(root_collection, Collection):
             raise TypeError(
-                f"root_collection must be an instance of Collection, not {type(root_collection)}"
+                f"root_collection must be an instance of Collection[Any], not {type(root_collection)}"
             )
         _disallow_transactions(session)
 
@@ -1817,7 +1817,7 @@ class GridOutChunkIterator:
     def __init__(
         self,
         grid_out: GridOut,
-        chunks: Collection,
+        chunks: Collection[Any],
         session: Optional[ClientSession],
         next_chunk: Any,
     ) -> None:
@@ -1830,7 +1830,7 @@ def __init__(
         self._num_chunks = math.ceil(float(self._length) / self._chunk_size)
         self._cursor = None
 
-    _cursor: Optional[Cursor]
+    _cursor: Optional[Cursor[Any]]
 
     def expected_chunk_length(self, chunk_n: int) -> int:
         if chunk_n < self._num_chunks - 1:
@@ -1908,7 +1908,7 @@ def close(self) -> None:
 
 
 class GridOutIterator:
-    def __init__(self, grid_out: GridOut, chunks: Collection, session: ClientSession):
+    def __init__(self, grid_out: GridOut, chunks: Collection[Any], session: ClientSession):
         self._chunk_iter = GridOutChunkIterator(grid_out, chunks, session, 0)
 
     def __iter__(self) -> GridOutIterator:
@@ -1921,14 +1921,14 @@ def next(self) -> bytes:
     __next__ = next
 
 
-class GridOutCursor(Cursor):
+class GridOutCursor(Cursor):  # type: ignore[type-arg]
     """A cursor / iterator for returning GridOut objects as the result
     of an arbitrary query against the GridFS files collection.
     """
 
     def __init__(
         self,
-        collection: Collection,
+        collection: Collection[Any],
         filter: Optional[Mapping[str, Any]] = None,
         skip: int = 0,
         limit: int = 0,

pymongo/asynchronous/auth_oidc.py

@@ -259,7 +259,7 @@ async def _sasl_continue_jwt(
     ) -> Mapping[str, Any]:
         self.access_token = None
         self.refresh_token = None
-        start_payload: dict = bson.decode(start_resp["payload"])
+        start_payload: dict[str, Any] = bson.decode(start_resp["payload"])
         if "issuer" in start_payload:
             self.idp_info = OIDCIdPInfo(**start_payload)
         access_token = await self._get_access_token()

pymongo/asynchronous/bulk.py

@@ -248,7 +248,7 @@ async def write_command(
         request_id: int,
         msg: bytes,
         docs: list[Mapping[str, Any]],
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
     ) -> dict[str, Any]:
         """A proxy for SocketInfo.write_command that handles event publishing."""
         cmd[bwc.field] = docs
@@ -334,7 +334,7 @@ async def unack_write(
         msg: bytes,
         max_doc_size: int,
         docs: list[Mapping[str, Any]],
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
     ) -> Optional[Mapping[str, Any]]:
         """A proxy for AsyncConnection.unack_write that handles event publishing."""
         if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):
@@ -419,7 +419,7 @@ async def _execute_batch_unack(
         bwc: Union[_BulkWriteContext, _EncryptedBulkWriteContext],
         cmd: dict[str, Any],
         ops: list[Mapping[str, Any]],
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
     ) -> list[Mapping[str, Any]]:
         if self.is_encrypted:
             _, batched_cmd, to_send = bwc.batch_command(cmd, ops)
@@ -446,7 +446,7 @@ async def _execute_batch(
         bwc: Union[_BulkWriteContext, _EncryptedBulkWriteContext],
         cmd: dict[str, Any],
         ops: list[Mapping[str, Any]],
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
     ) -> tuple[dict[str, Any], list[Mapping[str, Any]]]:
         if self.is_encrypted:
             _, batched_cmd, to_send = bwc.batch_command(cmd, ops)

pymongo/asynchronous/change_stream.py

@@ -164,7 +164,7 @@ def _aggregation_command_class(self) -> Type[_AggregationCommand]:
         raise NotImplementedError
 
     @property
-    def _client(self) -> AsyncMongoClient:
+    def _client(self) -> AsyncMongoClient:  # type: ignore[type-arg]
         """The client against which the aggregation commands for
         this AsyncChangeStream will be run.
         """
@@ -206,7 +206,7 @@ def _command_options(self) -> dict[str, Any]:
     def _aggregation_pipeline(self) -> list[dict[str, Any]]:
         """Return the full aggregation pipeline for this AsyncChangeStream."""
         options = self._change_stream_options()
-        full_pipeline: list = [{"$changeStream": options}]
+        full_pipeline: list[dict[str, Any]] = [{"$changeStream": options}]
         full_pipeline.extend(self._pipeline)
         return full_pipeline
 
@@ -237,7 +237,7 @@ def _process_result(self, result: Mapping[str, Any], conn: AsyncConnection) -> N
 
     async def _run_aggregation_cmd(
         self, session: Optional[AsyncClientSession], explicit_session: bool
-    ) -> AsyncCommandCursor:
+    ) -> AsyncCommandCursor:  # type: ignore[type-arg]
         """Run the full aggregation pipeline for this AsyncChangeStream and return
         the corresponding AsyncCommandCursor.
         """
@@ -257,7 +257,7 @@ async def _run_aggregation_cmd(
             operation=_Op.AGGREGATE,
         )
 
-    async def _create_cursor(self) -> AsyncCommandCursor:
+    async def _create_cursor(self) -> AsyncCommandCursor:  # type: ignore[type-arg]
         async with self._client._tmp_session(self._session, close=False) as s:
             return await self._run_aggregation_cmd(
                 session=s, explicit_session=self._session is not None

pymongo/asynchronous/client_bulk.py

@@ -88,7 +88,7 @@ class _AsyncClientBulk:
 
     def __init__(
         self,
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
         write_concern: WriteConcern,
         ordered: bool = True,
         bypass_document_validation: Optional[bool] = None,
@@ -233,7 +233,7 @@ async def write_command(
         msg: Union[bytes, dict[str, Any]],
         op_docs: list[Mapping[str, Any]],
         ns_docs: list[Mapping[str, Any]],
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
     ) -> dict[str, Any]:
         """A proxy for AsyncConnection.write_command that handles event publishing."""
         cmd["ops"] = op_docs
@@ -324,7 +324,7 @@ async def unack_write(
         msg: bytes,
         op_docs: list[Mapping[str, Any]],
         ns_docs: list[Mapping[str, Any]],
-        client: AsyncMongoClient,
+        client: AsyncMongoClient[Any],
     ) -> Optional[Mapping[str, Any]]:
         """A proxy for AsyncConnection.unack_write that handles event publishing."""
         if _COMMAND_LOGGER.isEnabledFor(logging.DEBUG):