Bug fix and improvement

aggify/aggify.py

@@ -12,6 +12,7 @@
     InvalidEmbeddedField,
     OutStageError,
     InvalidArgument,
+    InvalidProjection,
 )
 from aggify.types import QueryParams, CollectionType
 from aggify.utilty import (
@@ -84,6 +85,9 @@ def project(self, **kwargs: QueryParams) -> "Aggify":
             Aggify: Returns an instance of the Aggify class for potential method chaining.
         """
 
+        if all([i in kwargs.values() for i in [0, 1]]):
+            raise InvalidProjection()
+
         # Extract fields to keep and check if _id should be deleted
         to_keep_values = {"id"}
         delete_id = kwargs.get("id") is not None
@@ -99,14 +103,16 @@ def project(self, **kwargs: QueryParams) -> "Aggify":
                 to_keep_values.add(key)
                 self.base_model._fields[key] = mongoengine_fields.IntField()  # noqa
             projection[get_db_field(self.base_model, key)] = value  # noqa
+            if value == 0:
+                del self.base_model._fields[key]  # noqa
 
         # Remove fields from the base model, except the ones in to_keep_values and possibly _id
-        keys_for_deletion = self.base_model._fields.keys() - to_keep_values  # noqa
-        if delete_id:
-            keys_for_deletion.add("id")
-        for key in keys_for_deletion:
-            del self.base_model._fields[key]  # noqa
-
+        if to_keep_values != {"id"}:
+            keys_for_deletion = self.base_model._fields.keys() - to_keep_values  # noqa
+            if delete_id:
+                keys_for_deletion.add("id")
+            for key in keys_for_deletion:
+                del self.base_model._fields[key]  # noqa
         # Append the projection stage to the pipelines
         self.pipelines.append({"$project": projection})
 
@@ -157,10 +163,14 @@ def add_fields(self, **fields) -> "Aggify":  # noqa
                 add_fields_stage["$addFields"][field] = {"$literal": expression}
             elif isinstance(expression, F):
                 add_fields_stage["$addFields"][field] = expression.to_dict()
-            elif isinstance(expression, list):
+            elif isinstance(expression, (list, dict)):
                 add_fields_stage["$addFields"][field] = expression
             elif isinstance(expression, Cond):
                 add_fields_stage["$addFields"][field] = dict(expression)
+            elif isinstance(expression, Q):
+                add_fields_stage["$addFields"][field] = convert_match_query(
+                    dict(expression)
+                )["$match"]
             else:
                 raise AggifyValueError([str, F, list], type(expression))
             # TODO: Should be checked if new field is embedded, create embedded field.

aggify/exceptions.py

@@ -74,3 +74,9 @@ def __init__(self, expected_list: list):
         self.message = f"Input is not correctly passed, expected {[expected for expected in expected_list]}"
         self.expecteds = expected_list
         super().__init__(self.message)
+
+
+class InvalidProjection(AggifyBaseException):
+    def __init__(self):
+        self.message = "You can't use inclusion and exclusion together."
+        super().__init__(self.message)

tests/test_aggify.py

@@ -12,6 +12,7 @@
     AlreadyExistsField,
     InvalidEmbeddedField,
     MongoIndexError,
+    InvalidProjection,
 )
 
 
@@ -53,9 +54,9 @@ def test_filtering_and_projection(self):
 
     def test_filtering_and_projection_with_deleting_id(self):
         aggify = Aggify(BaseModel)
-        aggify.filter(age__gte=30).project(name=1, age=1, id=0)
+        aggify.filter(age__gte=30).project(name=1, age=1, id=1)
         assert len(aggify.pipelines) == 2
-        assert aggify.pipelines[1]["$project"] == {"_id": 0, "name": 1, "age": 1}
+        assert aggify.pipelines[1]["$project"] == {"_id": 1, "name": 1, "age": 1}
 
     def test_filtering_and_ordering(self):
         aggify = Aggify(BaseModel)
@@ -617,3 +618,21 @@ def test_aggify_get_item_slice_negative_start(self):
     def test_group_invalid_field(self):
         thing = list(Aggify(BaseModel).group("invalid").annotate("name", "first", 2))
         assert thing[0]["$group"] == {"_id": "invalid", "name": {"$first": 2}}
+
+    def test_add_field_with_q(self):
+        thing = list(Aggify(BaseModel).add_fields(new_field=Q(name__exact="aggify")))
+        assert thing[0]["$addFields"] == {"new_field": {"$eq": ["$name", "aggify"]}}
+
+    def test_add_field_with_dict(self):
+        thing = list(
+            Aggify(BaseModel).add_fields(new_field={"$eq": ["$owner.visibility", 0]})
+        )
+        assert thing[0]["$addFields"] == {
+            "new_field": {"$eq": ["$owner.visibility", 0]}
+        }
+
+    def test_project_use_inclusion_and_exclusion_together(self):
+        aggify = Aggify(BaseModel)
+        with pytest.raises(InvalidProjection):
+            # noinspection PyUnusedLocal
+            var = aggify.project(name=0, age=1)

tests/test_query.py

@@ -82,10 +82,10 @@ class ParameterTestCase:
     ParameterTestCase(
         compiled_query=Aggify(PostDocument)
         .filter(caption__contains="hello")
-        .project(caption=1, deleted_at=0),
+        .project(caption=1, deleted_at=1),
         expected_query=[
             {"$match": {"caption": {"$regex": "hello"}}},
-            {"$project": {"caption": 1, "deleted_at": 0}},
+            {"$project": {"caption": 1, "deleted_at": 1}},
         ],
     ),
     ParameterTestCase(