feat: add aggregation utility, tests, and some nitpicking here and there, refs #19

Author: volfpeter

motorhead/__init__.py

@@ -1,4 +1,7 @@
 from . import operator as operator
+from .aggregation import Aggregation as Aggregation
+from .aggregation import AggregationStage as AggregationStage
+from .aggregation import make_aggregation_stage as make_aggregation_stage
 from .bound_method_wrapper import BoundMethodWrapper as BoundMethodWrapper
 from .delete_rule import DeleteConfig as DeleteConfig
 from .delete_rule import DeleteError as DeleteError
@@ -13,13 +16,9 @@
 from .query import Q as Q
 from .query import Query as Query
 from .query import Queryable as Queryable
-from .service import DeleteResult as DeleteResult
-from .service import InsertManyResult as InsertManyResult
-from .service import InsertOneResult as InsertOneResult
 from .service import Service as Service
 from .service import ServiceConfig as ServiceConfig
 from .service import ServiceException as ServiceException
-from .service import UpdateResult as UpdateResult
 from .typing import AgnosticClient as AgnosticClient
 from .typing import AgnosticClientSession as AgnosticClientSession
 from .typing import AgnosticCollection as AgnosticCollection
@@ -35,15 +34,19 @@
 from .typing import CollectionOptions as CollectionOptions
 from .typing import DatabaseProvider as DatabaseProvider
 from .typing import DeleteOptions as DeleteOptions
+from .typing import DeleteResult as DeleteResult
 from .typing import FindOptions as FindOptions
 from .typing import IndexData as IndexData
 from .typing import InsertManyOptions as InsertManyOptions
+from .typing import InsertManyResult as InsertManyResult
 from .typing import InsertOneOptions as InsertOneOptions
+from .typing import InsertOneResult as InsertOneResult
 from .typing import MongoProjection as MongoProjection
 from .typing import MongoQuery as MongoQuery
 from .typing import UpdateManyOptions as UpdateManyOptions
 from .typing import UpdateObject as UpdateObject
 from .typing import UpdateOneOptions as UpdateOneOptions
+from .typing import UpdateResult as UpdateResult
 from .validator import ValidationError as ValidationError
 from .validator import Validator as Validator
 from .validator import validator as validator

motorhead/aggregation.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+    from typing import TypeAlias
+
+    from typing_extensions import Self
+
+    from .typing import Clause
+
+AggregationStage = Literal[
+    "$addFields",
+    "$bucket",
+    "$bucketAuto",
+    "$changeStream",
+    "$changeStreamSplitLargeEvent",
+    "$collStats",
+    "$count",
+    "$currentOp",
+    "$densify",
+    "$documents",
+    "$facet",
+    "$fill",
+    "$geoNear",
+    "$graphLookup",
+    "$group",
+    "$indexStats",
+    "$limit",
+    "$listLocalSessions",
+    "$listSampledQueries",
+    "$listSearchIndexes",
+    "$listSessions",
+    "$lookup",
+    "$match",
+    "$merge",
+    "$out",
+    "$planCacheStats",
+    "$project",
+    "$querySettings",
+    "$redact",
+    "$replaceRoot",
+    "$replaceWith",
+    "$sample",
+    "$search",
+    "$searchMeta",
+    "$set",
+    "$setWindowFields",
+    "$shardedDataDistribution",
+    "$skip",
+    "$sort",
+    "$sortByCount",
+    "$unionWith",
+    "$unset",
+    "$unwind",
+    "$vectorSearch",
+]
+"""Aggregation pipeline stage."""
+
+
+AggregationData: TypeAlias = Any
+
+
+def make_aggregation_stage(
+    stage: AggregationStage, value: AggregationData | Clause
+) -> dict[str, AggregationData]:
+    """
+    Creates an aggregation pipeline stage.
+
+    Arguments:
+        stage: The stage operator.
+        value: The stage operator's content.
+
+    Returns:
+        The aggregation pipeline stage.
+    """
+    return {stage: value.to_mongo() if hasattr(value, "to_mongo") else value}
+
+
+class Aggregation(list[dict[str, AggregationData]]):
+    """Aggregation pipeline."""
+
+    def __init__(self, stages: Iterable[AggregationData] = ()) -> None:
+        """
+        Initialization.
+
+        Arguments:
+            stages: The aggregation pipeline stages.
+        """
+        super().__init__(stages)
+
+    def stage(self, stage: AggregationStage, value: AggregationData | Clause) -> Self:
+        """
+        Adds the stage to the aggregation pipeline.
+
+        Arguments:
+            stage: The stage operator.
+            value: The stage operator's content.
+
+        Returns:
+            The aggregation pipeline.
+        """
+        self.append(make_aggregation_stage(stage, value))
+        return self

motorhead/operator.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Generator
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from collections.abc import Generator
+    from typing import Any
+
     from .query import Field
     from .typing import Clause
 

motorhead/query.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 from pydantic import BaseModel
 
@@ -23,12 +23,14 @@
     Size,
     Type,
 )
-from .typing import Clause
-
-_T = TypeVar("_T", bound=BaseModel)
 
 if TYPE_CHECKING:
-    from .typing import MongoQuery
+    from typing import Any
+
+    from .typing import Clause, MongoQuery
+
+
+_T = TypeVar("_T", bound=BaseModel)
 
 
 class Field:

motorhead/service.py

@@ -1,17 +1,19 @@
 from __future__ import annotations
 
-from collections.abc import AsyncGenerator, Callable, Coroutine, Generator, Iterable, Mapping, Sequence
 from contextlib import AbstractAsyncContextManager, asynccontextmanager, nullcontext
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypedDict, TypeVar, get_args
+from typing import TYPE_CHECKING, Generic, TypedDict, TypeVar, get_args
 
 from bson import ObjectId
 from pydantic import BaseModel
-from pymongo.results import DeleteResult, InsertManyResult, InsertOneResult, UpdateResult
 
+from .delete_rule import DeleteRule
 from .operator import ensure_dict
-from .typing import ClauseOrMongoQuery
+from .validator import Validator
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Callable, Coroutine, Generator, Iterable, Mapping, Sequence
+    from typing import Any, ClassVar
+
     from .typing import (
         AgnosticClient,
         AgnosticClientSession,
@@ -20,29 +22,30 @@
         AgnosticCursor,
         AgnosticDatabase,
         AgnosticLatentCommandCursor,
+        ClauseOrMongoQuery,
         Collation,
         CollectionOptions,
         DeleteOptions,
+        DeleteResult,
         FindOptions,
         IndexData,
         InsertManyOptions,
+        InsertManyResult,
         InsertOneOptions,
+        InsertOneResult,
         MongoProjection,
         UpdateManyOptions,
         UpdateObject,
         UpdateOneOptions,
+        UpdateResult,
     )
 
-from .delete_rule import DeleteRule
-from .validator import Validator
 
 __all__ = (
     "BaseService",
-    "DeleteResult",
-    "InsertManyResult",
-    "InsertOneResult",
     "Service",
-    "UpdateResult",
+    "ServiceConfig",
+    "ServiceException",
 )
 
 TInsert = TypeVar("TInsert", bound=BaseModel)

motorhead/typing.py

@@ -12,6 +12,10 @@
 from motor.core import AgnosticDatabase as _AgnosticDatabase
 from motor.core import AgnosticLatentCommandCursor as _AgnosticLatentCommandCursor
 from pymongo.collation import Collation as PMCollation
+from pymongo.results import DeleteResult as DeleteResult
+from pymongo.results import InsertManyResult as InsertManyResult
+from pymongo.results import InsertOneResult as InsertOneResult
+from pymongo.results import UpdateResult as UpdateResult
 
 if TYPE_CHECKING:
     from bson.codec_options import CodecOptions

motorhead/validator.py

@@ -1,9 +1,13 @@
-from collections.abc import Callable, Coroutine
-from typing import Literal, TypeVar
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, TypeVar
 
 from .bound_method_wrapper import BoundMethodWrapper
 from .typing import ClauseOrMongoQuery
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine
+
 __all__ = (
     "ValidationError",
     "Validator",

pyproject.toml

@@ -52,6 +52,7 @@ pytest = "^8.3.3"
 pytest-asyncio = "^0.24.0"
 pytest-docker = "^3.1.1"
 pytest-random-order = "^1.1.1"
+typing-extensions = "^4.12.2"
 
 [tool.mypy]
 strict = true