Skip to content

Commit d031820

Browse files
committed
Add docstring and refactor type checkers.
1 parent a6ca424 commit d031820

File tree

1 file changed

+109
-26
lines changed

1 file changed

+109
-26
lines changed

django_mongodb/fields/json.py

Lines changed: 109 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,22 +25,47 @@ def data_contains(self, compiler, connection): # noqa: ARG001
2525
raise NotSupportedError("contains lookup is not supported on this database backend.")
2626

2727

28+
def _has_key_predicate(path, root_column, negated=False):
29+
result = {"$and": [{"$ne": [{"$type": path}, "missing"]}, {"$ne": [root_column, None]}]}
30+
if negated:
31+
result = {"$not": result}
32+
return result
33+
34+
2835
def has_key_lookup(self, compiler, connection):
36+
"""
37+
Performs a key lookup in a JSONField to check for the existence of a key.
38+
39+
The key lookup in a JSONField is defined by the following conditions:
40+
1. The path type must be different from "missing".
41+
2. In cases where the JSONField is None (as the $type returns Null instead of "missing"),
42+
an additional check for not null is added.
43+
44+
Returns:
45+
A dictionary representing the MongoDB query for key lookup.
46+
47+
Raises:
48+
AssertionError: If `self.mongo_operator` is None and there are multiple keys.
49+
50+
Note:
51+
The function handles key lookup in two ways:
52+
- Via `KeyTransform`
53+
- Directly via key
54+
55+
Both are transformed into a `KeyTransform` for consistent handling.
56+
"""
2957
rhs = self.rhs
3058
lhs = process_lhs(self, compiler, connection)
31-
if not isinstance(rhs, (list | tuple)):
59+
if not isinstance(rhs, list | tuple):
3260
rhs = [rhs]
3361
paths = []
62+
# Transform all keys into KeyTransform instances for consistent handling
3463
for key in rhs:
35-
if isinstance(key, KeyTransform):
36-
rhs_json_path = key.as_mql(compiler, connection)
37-
else:
38-
rhs_json_path = KeyTransform(key, self.lhs).as_mql(compiler, connection)
39-
paths.append(rhs_json_path)
40-
64+
rhs_json_path = key if isinstance(key, KeyTransform) else KeyTransform(key, self.lhs)
65+
paths.append(rhs_json_path.as_mql(compiler, connection))
4166
keys = []
4267
for path in paths:
43-
keys.append({"$and": [{"$ne": [{"$type": path}, "missing"]}, {"$ne": [lhs, None]}]})
68+
keys.append(_has_key_predicate(path, lhs))
4469
if self.mongo_operator is None:
4570
assert len(keys) == 1
4671
return keys[0]
@@ -60,15 +85,34 @@ def json_exact_process_rhs(self, compiler, connection):
6085

6186

6287
def key_transform(self, compiler, connection):
88+
"""
89+
Transforms a Django KeyTransform (JSON path) into a MongoDB query path.
90+
91+
In MongoDB, JSON paths cannot always be represented simply as $var.key1.key2.key3
92+
due to potential array types. Therefore, indexing arrays requires the use of
93+
`arrayElemAt`. Additionally, a conditional check (if statement) is necessary
94+
to verify the type before performing the operation.
95+
96+
Returns:
97+
A dictionary representing the MongoDB query for the JSON path.
98+
99+
Note:
100+
The function constructs the MongoDB path by iterating through the key transforms
101+
and handling both field access and array indexing appropriately.
102+
"""
63103
key_transforms = [self.key_name]
64104
previous = self.lhs
105+
# Collect all key transforms in order.
65106
while isinstance(previous, KeyTransform):
66107
key_transforms.insert(0, previous.key_name)
67108
previous = previous.lhs
68109
lhs_mql = previous.as_mql(compiler, connection)
69110
result = lhs_mql
111+
# Build the MongoDB path using the collected key transforms
70112
for key in key_transforms:
71113
get_field = {"$getField": {"input": result, "field": key}}
114+
# Handle array indexing if the key is a digit
115+
# if we have an index like '001' is not an array index.
72116
if key.isdigit() and str(int(key)) == key:
73117
result = {
74118
"$cond": {
@@ -82,46 +126,85 @@ def key_transform(self, compiler, connection):
82126
return result
83127

84128

85-
def _expr_type_in(mql, types, negation=False):
86-
result = {"$in": [{"$type": mql}, types]}
87-
if negation:
88-
result = {"$not": result}
89-
return result
129+
def key_transform_in(self, compiler, connection):
130+
"""
131+
Checks if a JSON path's values are within a set of specified values and ensures the key exists.
90132
133+
This function performs two main checks:
134+
1. Verifies that the resulting path values are within the right-hand side (rhs) values.
135+
2. Ensures that the key exists.
91136
92-
def key_transform_in(self, compiler, connection):
137+
Returns:
138+
A dictionary representing the MongoDB query that checks both conditions.
139+
140+
Note:
141+
The function processes the left-hand side (lhs) and right-hand side (rhs)
142+
of the query, constructs the MongoDB expression, and checks both the value
143+
existence in rhs and the key existence in the JSON document.
144+
"""
93145
lhs_mql = process_lhs(self, compiler, connection)
146+
# Traverse to the root column
147+
previous = self.lhs
148+
while isinstance(previous, KeyTransform):
149+
previous = previous.lhs
150+
root_column = previous.as_mql(compiler, connection)
94151
value = process_rhs(self, compiler, connection)
152+
# Construct the expression to check if lhs_mql values are in rhs values
95153
expr = connection.mongo_operators[self.lookup_name](lhs_mql, value)
96-
type_in = _expr_type_in(lhs_mql, ["missing", "null"], True)
97-
return {"$and": [expr, type_in]}
154+
return {"$and": [_has_key_predicate(lhs_mql, root_column), expr]}
98155

99156

100157
def key_transform_is_null(self, compiler, connection):
101158
"""
102-
The KeyTransformIsNull lookup borrows the logic from HasKey for isnull=False.
103-
If isnull=True, the query should only match objects that don't have the key.
104-
https://code.djangoproject.com/ticket/32252
159+
Handles the KeyTransformIsNull lookup.
160+
161+
This function borrows logic from HasKey for `isnull=False`. If `isnull=True`,
162+
the query should match only objects that do not have the key.
163+
164+
Returns:
165+
A dictionary representing the MongoDB query for checking the nullability of the key.
166+
167+
Note:
168+
If `isnull=True`, the query matches objects where the key is missing
169+
or the root column is null.
170+
If `isnull=False`, the query negates the result to match objects where the key exists
171+
and the root column isn't null.
172+
173+
Reference:
174+
https://code.djangoproject.com/ticket/32252
105175
"""
106176
lhs_mql = process_lhs(self, compiler, connection)
107177
rhs_mql = process_rhs(self, compiler, connection)
178+
108179
# Get the root column.
109180
previous = self.lhs
110181
while isinstance(previous, KeyTransform):
111182
previous = previous.lhs
112183
root_column = previous.as_mql(compiler, connection)
113-
type_in = _expr_type_in(lhs_mql, ["missing"])
114-
result = {"$or": [type_in, {"$eq": [root_column, None]}]}
115-
if not rhs_mql:
116-
result = {"$not": result}
117-
return result
184+
return _has_key_predicate(lhs_mql, root_column, negated=rhs_mql)
118185

119186

120187
def key_transform_numeric_lookup_mixin(self, compiler, connection):
188+
"""
189+
Checks if the field exists and fulfills the given numeric lookup expression.
190+
191+
This function ensures that the field in the JSON document:
192+
1. Exists (i.e., is not "missing" or "null").
193+
2. Satisfies the specified numeric lookup expression.
194+
195+
Returns:
196+
A dictionary representing the MongoDB query that checks both conditions.
197+
198+
Note:
199+
The function constructs the MongoDB expression to check if the field exists
200+
and fulfills the numeric lookup expression, and combines these conditions
201+
using $and.
202+
"""
121203
expr = builtin_lookup(self, compiler, connection)
122204
lhs = process_lhs(self, compiler, connection)
123-
type_in = _expr_type_in(lhs, ["missing", "null"], True)
124-
return {"$and": [expr, type_in]}
205+
# Check if the type of lhs is not "missing" or "null"
206+
not_missing_or_null = {"$not": {"$in": [{"$type": lhs}, ["missing", "null"]]}}
207+
return {"$and": [expr, not_missing_or_null]}
125208

126209

127210
def register_json_field():

0 commit comments

Comments
 (0)