mongo-aggro 0.1.0__py3-none-any.whl

mongo_aggro/base.py

@@ -0,0 +1,195 @@
+"""Base classes for MongoDB aggregation pipeline stages."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterator
+from typing import Any, Self
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import CoreSchema, core_schema
+
+# Sort direction constants for use with Sort stage and with_sort method
+ASCENDING: int = 1
+DESCENDING: int = -1
+
+# Type alias for sort specification used in aggregation
+# Uses int to be compatible with both Literal[-1, 1] and regular int values
+SortSpec = dict[str, int]
+# Type alias for aggregation input tuple (pipeline, sort)
+AggregationInput = tuple[list[dict[str, Any]], SortSpec]
+
+
+class BaseStage(ABC):
+    """Abstract base class for all MongoDB aggregation pipeline stages."""
+
+    @abstractmethod
+    def model_dump(self) -> dict[str, Any]:
+        """
+        Convert the stage to its MongoDB dictionary representation.
+
+        Returns:
+            dict[str, Any]: MongoDB aggregation stage dictionary
+        """
+        pass
+
+
+class Pipeline:
+    """
+    MongoDB aggregation pipeline builder.
+
+    This class acts as a container for aggregation stages and can be
+    directly passed to MongoDB's aggregate() method. It implements
+    __iter__ to allow MongoDB drivers to iterate through the stages.
+
+    Example:
+        >>> pipeline = Pipeline()
+        >>> pipeline.add_stage(Match(field="status", value="active"))
+        >>> pipeline.add_stage(Group(id="$category", count={"$sum": 1}))
+        >>> collection.aggregate(pipeline)
+
+        >>> # Or with constructor
+        >>> pipeline = Pipeline([
+        ...     Match(field="status", value="active"),
+        ...     Unwind(path="items")
+        ... ])
+    """
+
+    def __init__(self, stages: list[BaseStage] | None = None) -> None:
+        """
+        Initialize the pipeline with optional initial stages.
+
+        Args:
+            stages: Optional list of initial pipeline stages
+        """
+        self._stages: list[BaseStage] = stages or []
+
+    def add_stage(self, stage: BaseStage) -> Self:
+        """
+        Append a new stage to the pipeline.
+
+        Args:
+            stage: A pipeline stage instance
+
+        Returns:
+            Self: Self for method chaining
+        """
+        self._stages.append(stage)
+        return self
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        """
+        Iterate through pipeline stages as dictionaries.
+
+        This allows the pipeline to be directly passed to MongoDB's
+        aggregate() method without calling any additional methods.
+
+        Yields:
+            dict[str, Any]: Each stage's MongoDB dictionary representation
+        """
+        for stage in self._stages:
+            yield stage.model_dump()
+
+    def __len__(self) -> int:
+        """Return the number of stages in the pipeline."""
+        return len(self._stages)
+
+    def __getitem__(self, index: int) -> BaseStage:
+        """Get a stage by index."""
+        return self._stages[index]
+
+    def to_list(self) -> list[dict[str, Any]]:
+        """
+        Convert the entire pipeline to a list of dictionaries.
+
+        Returns:
+            list[dict[str, Any]]: List of MongoDB stage dictionaries
+        """
+        return list(self)
+
+    def with_sort(self, sort: SortSpec) -> AggregationInput:
+        """
+        Return pipeline as aggregation input tuple with sort specification.
+
+        This is useful for integrating with pagination utilities that
+        expect a tuple of (pipeline, sort_spec). Common in Beanie ODM
+        pagination patterns.
+
+        Args:
+            sort: Sort specification dict with field names as keys
+                  and -1 (descending) or 1 (ascending) as values
+
+        Returns:
+            AggregationInput: Tuple of (pipeline_list, sort_spec)
+
+        Example:
+            >>> pipeline = Pipeline([
+            ...     Match(query={"status": "active"}),
+            ...     Group(id="$category", total={"$sum": "$amount"})
+            ... ])
+            >>> # Use with pagination utilities
+            >>> aggregation_input = pipeline.with_sort({"total": -1})
+            >>> await query.paginate_and_cast(
+            ...     query=base_query,
+            ...     projection_model=OutputModel,
+            ...     aggregation_input=aggregation_input
+            ... )
+        """
+        return (self.to_list(), sort)
+
+    def extend(self, stages: list[BaseStage]) -> Self:
+        """
+        Extend the pipeline with multiple stages.
+
+        Args:
+            stages: List of pipeline stages to add
+
+        Returns:
+            Self: Self for method chaining
+        """
+        self._stages.extend(stages)
+        return self
+
+    def extend_raw(self, raw_stages: list[dict[str, Any]]) -> Self:
+        """
+        Extend the pipeline with raw dictionary stages.
+
+        This is useful when combining typed stages with raw MongoDB
+        pipeline stages that may not have a corresponding class.
+
+        Args:
+            raw_stages: List of raw MongoDB stage dictionaries
+
+        Returns:
+            Self: Self for method chaining
+
+        Example:
+            >>> pipeline = Pipeline([Match(query={"status": "active"})])
+            >>> pipeline.extend_raw([
+            ...     {"$addFields": {"computed": {"$multiply": ["$a", "$b"]}}},
+            ...     {"$project": {"_id": 0, "result": "$computed"}}
+            ... ])
+        """
+        for raw_stage in raw_stages:
+            self._stages.append(_RawStage(raw_stage))
+        return self
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Any, handler: GetCoreSchemaHandler
+    ) -> CoreSchema:
+        """
+        Allow Pipeline to be used as a Pydantic field type.
+
+        This enables using Pipeline in other Pydantic models,
+        for example in Lookup's pipeline field.
+        """
+        return core_schema.is_instance_schema(cls)
+
+
+class _RawStage(BaseStage):
+    """Internal class for wrapping raw dictionary stages."""
+
+    def __init__(self, raw: dict[str, Any]) -> None:
+        self._raw = raw
+
+    def model_dump(self) -> dict[str, Any]:
+        return self._raw

mongo_aggro/operators.py

@@ -0,0 +1,247 @@
+"""Query and expression operators for MongoDB aggregation."""
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class QueryOperator(BaseModel):
+    """Base class for query operators used in $match and other stages."""
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        extra="forbid",
+    )
+
+
+class And(QueryOperator):
+    """
+    Logical AND operator for combining multiple conditions.
+
+    Example:
+        >>> And(conditions=[
+        ...     {"status": "active"},
+        ...     {"age": {"$gt": 18}}
+        ... ]).model_dump()
+        {"$and": [{"status": "active"}, {"age": {"$gt": 18}}]}
+    """
+
+    conditions: list[dict[str, Any]] = Field(
+        ..., description="List of conditions to AND together"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$and": self.conditions}
+
+
+class Or(QueryOperator):
+    """
+    Logical OR operator for combining multiple conditions.
+
+    Example:
+        >>> Or(conditions=[
+        ...     {"status": "active"},
+        ...     {"status": "pending"}
+        ... ]).model_dump()
+        {"$or": [{"status": "active"}, {"status": "pending"}]}
+    """
+
+    conditions: list[dict[str, Any]] = Field(
+        ..., description="List of conditions to OR together"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$or": self.conditions}
+
+
+class Not(QueryOperator):
+    """
+    Logical NOT operator for negating a condition.
+
+    Example:
+        >>> Not(condition={"$regex": "^test"}).model_dump()
+        {"$not": {"$regex": "^test"}}
+    """
+
+    condition: dict[str, Any] = Field(..., description="Condition to negate")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$not": self.condition}
+
+
+class Nor(QueryOperator):
+    """
+    Logical NOR operator - matches documents that fail all conditions.
+
+    Example:
+        >>> Nor(conditions=[
+        ...     {"price": {"$gt": 1000}},
+        ...     {"rating": {"$lt": 3}}
+        ... ]).model_dump()
+        {"$nor": [{"price": {"$gt": 1000}}, {"rating": {"$lt": 3}}]}
+    """
+
+    conditions: list[dict[str, Any]] = Field(
+        ..., description="List of conditions to NOR together"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$nor": self.conditions}
+
+
+class Expr(QueryOperator):
+    """
+    $expr operator for using aggregation expressions in queries.
+
+    Example:
+        >>> Expr(expression={"$eq": ["$field1", "$field2"]}).model_dump()
+        {"$expr": {"$eq": ["$field1", "$field2"]}}
+    """
+
+    expression: dict[str, Any] = Field(
+        ..., description="Aggregation expression"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$expr": self.expression}
+
+
+# Comparison operators
+class Eq(QueryOperator):
+    """$eq comparison operator."""
+
+    value: Any = Field(..., description="Value to compare")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$eq": self.value}
+
+
+class Ne(QueryOperator):
+    """$ne (not equal) comparison operator."""
+
+    value: Any = Field(..., description="Value to compare")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$ne": self.value}
+
+
+class Gt(QueryOperator):
+    """$gt (greater than) comparison operator."""
+
+    value: Any = Field(..., description="Value to compare")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$gt": self.value}
+
+
+class Gte(QueryOperator):
+    """$gte (greater than or equal) comparison operator."""
+
+    value: Any = Field(..., description="Value to compare")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$gte": self.value}
+
+
+class Lt(QueryOperator):
+    """$lt (less than) comparison operator."""
+
+    value: Any = Field(..., description="Value to compare")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$lt": self.value}
+
+
+class Lte(QueryOperator):
+    """$lte (less than or equal) comparison operator."""
+
+    value: Any = Field(..., description="Value to compare")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$lte": self.value}
+
+
+class In(QueryOperator):
+    """$in operator - matches any value in the array."""
+
+    values: list[Any] = Field(..., description="List of values to match")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$in": self.values}
+
+
+class Nin(QueryOperator):
+    """$nin operator - matches none of the values in the array."""
+
+    values: list[Any] = Field(..., description="List of values to exclude")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$nin": self.values}
+
+
+class Regex(QueryOperator):
+    """$regex operator for pattern matching."""
+
+    pattern: str = Field(..., description="Regular expression pattern")
+    options: str | None = Field(
+        default=None, description="Regex options (i, m, x, s)"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        result: dict[str, Any] = {"$regex": self.pattern}
+        if self.options:
+            result["$options"] = self.options
+        return result
+
+
+class Exists(QueryOperator):
+    """$exists operator - matches documents where field exists/doesn't."""
+
+    exists: bool = Field(
+        default=True, description="True if field should exist"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$exists": self.exists}
+
+
+class Type(QueryOperator):
+    """$type operator - matches documents where field is of specified type."""
+
+    bson_type: str | int | list[str | int] = Field(
+        ..., description="BSON type(s) to match"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$type": self.bson_type}
+
+
+class ElemMatch(QueryOperator):
+    """$elemMatch operator - matches array elements."""
+
+    conditions: dict[str, Any] = Field(
+        ..., description="Conditions for array elements"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$elemMatch": self.conditions}
+
+
+class Size(QueryOperator):
+    """$size operator - matches arrays with specific length."""
+
+    size: int = Field(..., description="Array size to match")
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$size": self.size}
+
+
+class All(QueryOperator):
+    """$all operator - matches arrays containing all specified elements."""
+
+    values: list[Any] = Field(
+        ..., description="Values that must all be present"
+    )
+
+    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+        return {"$all": self.values}