handle aggregation method on the whole collection (#1203)

Author: CAPITAINMARVEL

beanie/odm/interfaces/aggregate.py

@@ -4,6 +4,7 @@
 from pydantic import BaseModel
 from pymongo.asynchronous.client_session import AsyncClientSession
 
+from beanie.odm.fields import ExpressionField
 from beanie.odm.queries.aggregation import AggregationQuery
 from beanie.odm.queries.find import FindMany
 
@@ -68,3 +69,110 @@ def aggregate(
             ignore_cache=ignore_cache,
             **pymongo_kwargs,
         )
+
+    @classmethod
+    async def sum(
+        cls,
+        field: Union[ExpressionField, float, int, str],
+        session: Optional[AsyncClientSession] = None,
+        ignore_cache: bool = False,
+    ) -> Optional[float]:
+        """
+        Sum of values of the given field over the entire collection.
+
+        Example:
+
+        ```python
+
+        class Sample(Document):
+            price: int
+
+        sum_count = await Document.sum(Sample.price)
+
+        ```
+
+        :param field: Union[ExpressionField, float, int, str]
+        :param session: Optional[AsyncClientSession] - pymongo session
+        :param ignore_cache: bool
+        :return: float - sum. None if there are no items.
+        """
+        return await cls.find_all().sum(field, session, ignore_cache)
+
+    @classmethod
+    async def avg(
+        cls,
+        field: Union[ExpressionField, float, int, str],
+        session: Optional[AsyncClientSession] = None,
+        ignore_cache: bool = False,
+    ) -> Optional[float]:
+        """
+        Average of values of the given field over the entire collection.
+
+        Example:
+
+        ```python
+
+        class Sample(Document):
+            price: int
+
+        avg_count = await Document.avg(Sample.price)
+        ```
+
+        :param field: Union[ExpressionField, float, int, str]
+        :param session: Optional[AsyncClientSession] - pymongo session
+        :param ignore_cache: bool
+        :return: Optional[float] - avg. None if there are no items.
+        """
+        return await cls.find_all().avg(field, session, ignore_cache)
+
+    @classmethod
+    async def max(
+        cls,
+        field: Union[ExpressionField, str, Any],
+        session: Optional[AsyncClientSession] = None,
+        ignore_cache: bool = False,
+    ) -> Optional[float]:
+        """
+        Max of the values of the given field over the entire collection.
+
+        Example:
+
+        ```python
+
+        class Sample(Document):
+            price: int
+
+        max_count = await Document.max(Sample.price)
+        ```
+
+        :param field: Union[ExpressionField, str, Any]
+        :param session: Optional[AsyncClientSession] - pymongo session
+        :return: float - max. None if there are no items.
+        """
+        return await cls.find_all().max(field, session, ignore_cache)
+
+    @classmethod
+    async def min(
+        cls,
+        field: Union[ExpressionField, str, Any],
+        session: Optional[AsyncClientSession] = None,
+        ignore_cache: bool = False,
+    ) -> Optional[float]:
+        """
+        Min of the values of the given field over the entire collection.
+
+        Example:
+
+        ```python
+
+        class Sample(Document):
+            price: int
+
+        min_count = await Document.min(Sample.price)
+        ```
+
+        :param field: Union[ExpressionField, str, Any]
+        :param session: Optional[AsyncClientSession] - pymongo session
+        :return: float - min. None if there are no items.
+        """
+        return await cls.find_all().min(field, session, ignore_cache)

beanie/odm/interfaces/aggregation_methods.py

@@ -22,7 +22,7 @@ def aggregate(
 
     async def sum(
         self,
-        field: Union[str, ExpressionField],
+        field: Union[ExpressionField, float, int, str],
         session: Optional[AsyncClientSession] = None,
         ignore_cache: bool = False,
     ) -> Optional[float]:
@@ -41,7 +41,7 @@ class Sample(Document):
 
         ```
 
-        :param field: Union[str, ExpressionField]
+        :param field: Union[ExpressionField, float, int, str]
         :param session: Optional[AsyncClientSession] - pymongo session
         :param ignore_cache: bool
         :return: float - sum. None if there are no items.
@@ -66,7 +66,7 @@ class Sample(Document):
 
     async def avg(
         self,
-        field,
+        field: Union[ExpressionField, float, int, str],
         session: Optional[AsyncClientSession] = None,
         ignore_cache: bool = False,
     ) -> Optional[float]:
@@ -84,7 +84,7 @@ class Sample(Document):
         avg_count = await Document.find(Sample.price <= 100).avg(Sample.count)
         ```
 
-        :param field: Union[str, ExpressionField]
+        :param field: Union[ExpressionField, float, int, str]
         :param session: Optional[AsyncClientSession] - pymongo session
         :param ignore_cache: bool
         :return: Optional[float] - avg. None if there are no items.
@@ -108,7 +108,7 @@ class Sample(Document):
 
     async def max(
         self,
-        field: Union[str, ExpressionField],
+        field: Union[ExpressionField, str, Any],
         session: Optional[AsyncClientSession] = None,
         ignore_cache: bool = False,
     ) -> Optional[float]:
@@ -126,7 +126,7 @@ class Sample(Document):
         max_count = await Document.find(Sample.price <= 100).max(Sample.count)
         ```
 
-        :param field: Union[str, ExpressionField]
+        :param field: Union[ExpressionField, str, Any]
         :param session: Optional[AsyncClientSession] - pymongo session
         :return: float - max. None if there are no items.
         """
@@ -149,7 +149,7 @@ class Sample(Document):
 
     async def min(
         self,
-        field: Union[str, ExpressionField],
+        field: Union[ExpressionField, str, Any],
         session: Optional[AsyncClientSession] = None,
         ignore_cache: bool = False,
     ) -> Optional[float]:
@@ -167,7 +167,7 @@ class Sample(Document):
         min_count = await Document.find(Sample.price <= 100).min(Sample.count)
         ```
 
-        :param field: Union[str, ExpressionField]
+        :param field: Union[ExpressionField, str, Any]
         :param session: Optional[AsyncClientSession] - pymongo session
         :return: float - min. None if there are no items.
         """

tests/odm/query/test_aggregate_methods.py

@@ -93,3 +93,81 @@ async def test_min_without_docs(session):
     )
 
     assert n is None
+
+
+async def test_all_sum(preset_documents, session):
+    n = await Sample.sum(Sample.increment)
+
+    assert n == 45
+
+    n = await Sample.sum(Sample.increment, session=session)
+
+    assert n == 45
+
+
+async def test_all_sum_without_docs(session):
+    n = await Sample.sum(Sample.increment)
+
+    assert n is None
+
+    n = await Sample.sum(Sample.increment, session=session)
+
+    assert n is None
+
+
+async def test_all_avg(preset_documents, session):
+    n = await Sample.avg(Sample.increment)
+
+    assert n == 4.5
+    n = await Sample.avg(Sample.increment, session=session)
+
+    assert n == 4.5
+
+
+async def test_all_avg_without_docs(session):
+    n = await Sample.avg(Sample.increment)
+
+    assert n is None
+    n = await Sample.avg(Sample.increment, session=session)
+
+    assert n is None
+
+
+async def test_all_max(preset_documents, session):
+    n = await Sample.max(Sample.increment)
+
+    assert n == 9
+
+    n = await Sample.max(Sample.increment, session=session)
+
+    assert n == 9
+
+
+async def test_all_max_without_docs(session):
+    n = await Sample.max(Sample.increment)
+
+    assert n is None
+
+    n = await Sample.max(Sample.increment, session=session)
+
+    assert n is None
+
+
+async def test_all_min(preset_documents, session):
+    n = await Sample.min(Sample.increment)
+
+    assert n == 0
+
+    n = await Sample.min(Sample.increment, session=session)
+
+    assert n == 0
+
+
+async def test_all_min_without_docs(session):
+    n = await Sample.min(Sample.increment)
+
+    assert n is None
+
+    n = await Sample.min(Sample.increment, session=session)
+
+    assert n is None