Skip to content

Commit

Permalink
add the addFields feature
Browse files Browse the repository at this point in the history
  • Loading branch information
seyed-dev committed Oct 28, 2023
1 parent 59e9c6c commit e64b807
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 1 deletion.
15 changes: 14 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand Down
23 changes: 23 additions & 0 deletions aggify/aggify.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
33 changes: 33 additions & 0 deletions tests/test_aggify.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit e64b807

Please sign in to comment.