fix wrapper usage

Author: blink1073

flask_pymongo/__init__.py

@@ -33,13 +33,13 @@
 import pymongo
 from flask import Flask, Response, abort, current_app, request
 from gridfs import GridFS, NoFile
-from pymongo import MongoClient, uri_parser
-from pymongo.database import Database
+from pymongo import uri_parser
 from pymongo.driver_info import DriverInfo
 from werkzeug.wsgi import wrap_file
 
 from flask_pymongo._version import __version__
 from flask_pymongo.helpers import BSONObjectIdConverter, BSONProvider
+from flask_pymongo.wrappers import Database, MongoClient
 
 DESCENDING = pymongo.DESCENDING
 """Descending sort order."""
@@ -64,8 +64,8 @@ class PyMongo:
     def __init__(
         self, app: Flask | None = None, uri: str | None = None, *args: Any, **kwargs: Any
     ) -> None:
-        self.cx: MongoClient[dict[str, Any]] | None = None
-        self.db: Database[dict[str, Any]] | None = None
+        self.cx: MongoClient | None = None
+        self.db: Database | None = None
 
         if app is not None:
             self.init_app(app, uri, *args, **kwargs)

flask_pymongo/wrappers.py

@@ -27,12 +27,70 @@
 from typing import Any
 
 from flask import abort
-from pymongo import collection
+from pymongo import collection, database, mongo_client
+
+
+class MongoClient(mongo_client.MongoClient[dict[str, Any]]):
+    """Wrapper for :class:`~pymongo.mongo_client.MongoClient`.
+
+    Returns instances of Flask-PyMongo
+    :class:`~flask_pymongo.wrappers.Database` instead of native PyMongo
+    :class:`~pymongo.database.Database` when accessed with dot notation.
+
+    """
+
+    def __getattr__(self, name: str) -> Any:
+        attr = super().__getattr__(name)
+        if isinstance(attr, database.Database):
+            return Database(self, name)
+        return attr
+
+    def __getitem__(self, name: str) -> Any:
+        attr = super().__getitem__(name)
+        if isinstance(attr, database.Database):
+            return Database(self, name)
+        return attr
+
+
+class Database(database.Database[dict[str, Any]]):
+    """Wrapper for :class:`~pymongo.database.Database`.
+
+    Returns instances of Flask-PyMongo
+    :class:`~flask_pymongo.wrappers.Collection` instead of native PyMongo
+    :class:`~pymongo.collection.Collection` when accessed with dot notation.
+
+    """
+
+    def __getattr__(self, name: str) -> Any:
+        attr = super().__getattr__(name)
+        if isinstance(attr, collection.Collection):
+            return Collection(self, name)
+        return attr
+
+    def __getitem__(self, name: str) -> Any:
+        item_ = super().__getitem__(name)
+        if isinstance(item_, collection.Collection):
+            return Collection(self, name)
+        return item_
 
 
 class Collection(collection.Collection[dict[str, Any]]):
     """Sub-class of PyMongo :class:`~pymongo.collection.Collection` with helpers."""
 
+    def __getattr__(self, name: str) -> Any:
+        attr = super().__getattr__(name)
+        if isinstance(attr, collection.Collection):
+            db = self._Collection__database
+            return Collection(db, attr.name)
+        return attr
+
+    def __getitem__(self, name: str) -> Any:
+        item = super().__getitem__(name)
+        if isinstance(item, collection.Collection):
+            db = self._Collection__database
+            return Collection(db, item.name)
+        return item
+
     def find_one_or_404(self, *args: Any, **kwargs: Any) -> Any:
         """Find a single document or raise a 404.