Merge pull request #10 from jorants/index-type-hints

Index type hints

beanie/__init__.py

@@ -1,6 +1,7 @@
 from beanie.cursor import Cursor
 from beanie.documents import Document
 from beanie.general import init_beanie
+from beanie.collection import Indexed
 
 __version__ = "0.3.2"
-__all__ = ["Document", "Cursor", "init_beanie"]
+__all__ = ["Document", "Cursor", "init_beanie", "Indexed"]

beanie/collection.py

@@ -2,7 +2,23 @@
 
 from motor.motor_asyncio import AsyncIOMotorCollection, AsyncIOMotorDatabase
 from pydantic.main import BaseModel
-from pymongo import IndexModel
+from pymongo import IndexModel, ASCENDING
+
+
+def Indexed(typ):
+    """
+    Returns a subclass of `typ` with an extra attribute `_indexed` et to True.
+    When instantiated the type of the result will actually be `typ`.
+    """
+
+    class NewType(typ):
+        _indexed = True
+
+        def __new__(cls, *args, **kwargs):
+            return typ.__new__(typ, *args, **kwargs)
+
+    NewType.__name__ = f"Indexed {typ.__name__}"
+    return NewType
 
 
 class IndexModelField(IndexModel):
@@ -55,14 +71,25 @@ async def collection_factory(
     # create motor collection
     collection = database[collection_parameters.name]
 
-    # create indexes
+    # Indexed field wrapped with Indexed()
+    found_indexes = [
+        IndexModel([(fname, ASCENDING)])
+        for fname, fvalue in document_class.__fields__.items()
+        if hasattr(fvalue.type_, "_indexed") and fvalue.type_._indexed
+    ]
+
+    # get indexes from the Collection class
     if collection_parameters.indexes:
-        await collection.create_indexes(collection_parameters.indexes)
+        found_indexes += collection_parameters.indexes
+
+    # Create indices
+    if found_indexes:
+        await collection.create_indexes(found_indexes)
 
     # create internal CollectionMeta class for the Document
     class CollectionMeta:
         name: str = collection_parameters.name
         motor_collection: AsyncIOMotorCollection = collection
-        indexes: List = collection_parameters.indexes
+        indexes: List = found_indexes
 
     return CollectionMeta

docs/index.md

@@ -311,7 +311,21 @@ class DocumentTestModelWithCustomCollectionName(Document):
 
 #### Indexes
 
-The indexes could be set up by the `indexes` field. It is a list where items could be:
+##### Simple indexes
+
+To setup an index over a single field the `Indexed` function can be used to wrap the type:
+
+```python
+from beanie import Indexed
+
+class DocumentTestModelWithIndex(Document):
+    test_int: Indexed(int)
+    test_list: List[SubDocument]
+    test_str: str
+```
+
+##### Complex indexes
+More complex indexes can be set up by the `indexes` field in a Collection class. It is a list where items could be:
 
 - single key. Name of the document's field
 - list of (key, direction) pairs. Key - string, name of the document's field. Direction - pymongo direction (
@@ -339,6 +353,8 @@ class DocumentTestModelWithIndex(Document):
         ]
 ```
 
+Complex and simple indices can be used in tandem.
+
 ### Use Motor Collection
 
 In case, when you need more low-level control, you can get access to the engine of the Beanie `Document`- [AsyncIO Motor Collection](https://motor.readthedocs.io/en/stable/api-asyncio/asyncio_motor_collection.html)
@@ -347,4 +363,4 @@ In case, when you need more low-level control, you can get access to the engine
 motor_collection = DocumentTestModel.get_motor_collection()
 await motor_collection.drop_index("index_name")
 ```
-
+

tests/models.py

@@ -4,7 +4,7 @@
 from pydantic import BaseModel
 from pymongo import IndexModel
 
-from beanie import Document
+from beanie import Document, Indexed
 
 
 class SubDocument(BaseModel):
@@ -28,6 +28,7 @@ class Collection:
 
 class DocumentTestModelWithIndex(Document):
     test_int: int
+    test_indexed_int: Indexed(int)
     test_list: List[SubDocument]
     test_str: str
 

tests/test_documents_init.py

@@ -33,6 +33,7 @@ async def test_index_creation():
     assert index_info == {
         "_id_": {"key": [("_id", 1)], "v": 2},
         "test_int_1": {"key": [("test_int", 1)], "v": 2},
+        "test_indexed_int_1": {"key": [("test_indexed_int", 1)], "v": 2},
         "test_int_1_test_str_-1": {
             "key": [("test_int", 1), ("test_str", -1)],
             "v": 2,