add the addFields feature

Author: seyed-dev

README.md

@@ -44,7 +44,7 @@ pip install aggify
 Here's a code snippet that demonstrates how to use Aggify to construct a MongoDB aggregation pipeline:
 
 ```python
-from aggify import Aggify, Q
+from aggify import Aggify, Q, F
 from mongoengine import Document, fields
 from pprint import pprint
 
@@ -131,6 +131,19 @@ pprint(
 # output:
 #         [{'$match': {'caption': 'hello'}}, {'$sort': {'_id': -1}}]
 
+pprint(
+    query.addFields({
+        "new_field_1": "some_string",
+        "new_field_2": F("existing_field") + 10,
+        "new_field_3": F("field_a") * F("field_b")}
+    ).pipelines
+)
+# output :
+#         [{'$addFields': {'new_field_1': {'$literal': 'some_string'},
+#                  'new_field_2': {'$add': ['$existing_field', 10]},
+#                  'new_field_3': {'$multiply': ['$field_a', '$field_b']}}}]
+
+
 ```
 
 In the sample usage above, you can see how Aggify simplifies the construction of MongoDB aggregation pipelines by allowing you to chain filters, projections, and other operations to build complex queries. The pprint(query.pipelines) line demonstrates how you can inspect the generated aggregation pipeline for debugging or analysis.

aggify/aggify.py

@@ -263,6 +263,29 @@ def raw(self, raw_query):
         self.pipelines.append(raw_query)
         return self
 
+    def addFields(self, fields):  # noqa
+        """
+        Generates a MongoDB addFields pipeline stage.
+
+        Args:
+            fields: A dictionary of field expressions and values.
+
+        Returns:
+            A MongoDB addFields pipeline stage.
+        """
+        add_fields_stage = {"$addFields": {}}
+
+        for field, expression in fields.items():
+            if isinstance(expression, str):
+                add_fields_stage["$addFields"][field] = {"$literal": expression}
+            elif isinstance(expression, F):
+                add_fields_stage["$addFields"][field] = expression.to_dict()
+            else:
+                raise ValueError("Invalid field expression")
+
+        self.pipelines.append(add_fields_stage)
+        return self
+
     def aggregate(self):
         """
         Returns the aggregated results.

tests/test_aggify.py

@@ -131,3 +131,36 @@ def test_filter_with_not_operator(self):
         aggify.filter(~Q(name="John"))
         assert len(aggify.pipelines) == 1
         assert aggify.pipelines[0]["$match"]["$not"][0]["name"] == "John"
+
+    def test_add_fields_string_literal(self):
+        aggify = Aggify(BaseModel)
+        fields = {
+            "new_field_1": "some_string",
+            "new_field_2": "another_string"
+        }
+        add_fields_stage = aggify.addFields(fields)
+
+        expected_stage = {
+            "$addFields": {
+                "new_field_1": {"$literal": "some_string"},
+                "new_field_2": {"$literal": "another_string"}
+            }
+        }
+
+        assert add_fields_stage.pipelines[0] == expected_stage
+
+    def test_add_fields_with_f_expression(self):
+        aggify = Aggify(BaseModel)
+        fields = {
+            "new_field_1": F("existing_field") + 10,
+            "new_field_2": F("field_a") * F("field_b")
+        }
+        add_fields_stage = aggify.addFields(fields)
+
+        expected_stage = {
+            "$addFields": {
+                "new_field_1": {"$add": ["$existing_field", 10]},
+                "new_field_2": {"$multiply": ["$field_a", "$field_b"]}
+            }
+        }
+        assert add_fields_stage.pipelines[0] == expected_stage