Add docstring and refactor type checkers.

Author: WaVEV

django_mongodb/fields/json.py

@@ -25,22 +25,47 @@ def data_contains(self, compiler, connection):  # noqa: ARG001
     raise NotSupportedError("contains lookup is not supported on this database backend.")
 
 
+def _has_key_predicate(path, root_column, negated=False):
+    result = {"$and": [{"$ne": [{"$type": path}, "missing"]}, {"$ne": [root_column, None]}]}
+    if negated:
+        result = {"$not": result}
+    return result
+
+
 def has_key_lookup(self, compiler, connection):
+    """
+    Performs a key lookup in a JSONField to check for the existence of a key.
+
+    The key lookup in a JSONField is defined by the following conditions:
+    1. The path type must be different from "missing".
+    2. In cases where the JSONField is None (as the $type returns Null instead of "missing"),
+       an additional check for not null is added.
+
+    Returns:
+        A dictionary representing the MongoDB query for key lookup.
+
+    Raises:
+        AssertionError: If `self.mongo_operator` is None and there are multiple keys.
+
+    Note:
+        The function handles key lookup in two ways:
+        - Via `KeyTransform`
+        - Directly via key
+
+        Both are transformed into a `KeyTransform` for consistent handling.
+    """
     rhs = self.rhs
     lhs = process_lhs(self, compiler, connection)
-    if not isinstance(rhs, (list | tuple)):
+    if not isinstance(rhs, list | tuple):
         rhs = [rhs]
     paths = []
+    # Transform all keys into KeyTransform instances for consistent handling
     for key in rhs:
-        if isinstance(key, KeyTransform):
-            rhs_json_path = key.as_mql(compiler, connection)
-        else:
-            rhs_json_path = KeyTransform(key, self.lhs).as_mql(compiler, connection)
-        paths.append(rhs_json_path)
-
+        rhs_json_path = key if isinstance(key, KeyTransform) else KeyTransform(key, self.lhs)
+        paths.append(rhs_json_path.as_mql(compiler, connection))
     keys = []
     for path in paths:
-        keys.append({"$and": [{"$ne": [{"$type": path}, "missing"]}, {"$ne": [lhs, None]}]})
+        keys.append(_has_key_predicate(path, lhs))
     if self.mongo_operator is None:
         assert len(keys) == 1
         return keys[0]
@@ -60,15 +85,34 @@ def json_exact_process_rhs(self, compiler, connection):
 
 
 def key_transform(self, compiler, connection):
+    """
+    Transforms a Django KeyTransform (JSON path) into a MongoDB query path.
+
+    In MongoDB, JSON paths cannot always be represented simply as $var.key1.key2.key3
+    due to potential array types. Therefore, indexing arrays requires the use of
+    `arrayElemAt`. Additionally, a conditional check (if statement) is necessary
+    to verify the type before performing the operation.
+
+    Returns:
+        A dictionary representing the MongoDB query for the JSON path.
+
+    Note:
+        The function constructs the MongoDB path by iterating through the key transforms
+        and handling both field access and array indexing appropriately.
+    """
     key_transforms = [self.key_name]
     previous = self.lhs
+    # Collect all key transforms in order.
     while isinstance(previous, KeyTransform):
         key_transforms.insert(0, previous.key_name)
         previous = previous.lhs
     lhs_mql = previous.as_mql(compiler, connection)
     result = lhs_mql
+    # Build the MongoDB path using the collected key transforms
     for key in key_transforms:
         get_field = {"$getField": {"input": result, "field": key}}
+        # Handle array indexing if the key is a digit
+        # if we have an index like '001' is not an array index.
         if key.isdigit() and str(int(key)) == key:
             result = {
                 "$cond": {
@@ -82,46 +126,85 @@ def key_transform(self, compiler, connection):
     return result
 
 
-def _expr_type_in(mql, types, negation=False):
-    result = {"$in": [{"$type": mql}, types]}
-    if negation:
-        result = {"$not": result}
-    return result
+def key_transform_in(self, compiler, connection):
+    """
+    Checks if a JSON path's values are within a set of specified values and ensures the key exists.
 
+    This function performs two main checks:
+    1. Verifies that the resulting path values are within the right-hand side (rhs) values.
+    2. Ensures that the key exists.
 
-def key_transform_in(self, compiler, connection):
+    Returns:
+        A dictionary representing the MongoDB query that checks both conditions.
+
+    Note:
+        The function processes the left-hand side (lhs) and right-hand side (rhs)
+        of the query, constructs the MongoDB expression, and checks both the value
+        existence in rhs and the key existence in the JSON document.
+    """
     lhs_mql = process_lhs(self, compiler, connection)
+    # Traverse to the root column
+    previous = self.lhs
+    while isinstance(previous, KeyTransform):
+        previous = previous.lhs
+    root_column = previous.as_mql(compiler, connection)
     value = process_rhs(self, compiler, connection)
+    # Construct the expression to check if lhs_mql values are in rhs values
     expr = connection.mongo_operators[self.lookup_name](lhs_mql, value)
-    type_in = _expr_type_in(lhs_mql, ["missing", "null"], True)
-    return {"$and": [expr, type_in]}
+    return {"$and": [_has_key_predicate(lhs_mql, root_column), expr]}
 
 
 def key_transform_is_null(self, compiler, connection):
     """
-    The KeyTransformIsNull lookup borrows the logic from HasKey for isnull=False.
-    If isnull=True, the query should only match objects that don't have the key.
-    https://code.djangoproject.com/ticket/32252
+    Handles the KeyTransformIsNull lookup.
+
+    This function borrows logic from HasKey for `isnull=False`. If `isnull=True`,
+    the query should match only objects that do not have the key.
+
+    Returns:
+        A dictionary representing the MongoDB query for checking the nullability of the key.
+
+    Note:
+        If `isnull=True`, the query matches objects where the key is missing
+                          or the root column is null.
+        If `isnull=False`, the query negates the result to match objects where the key exists
+                           and the root column isn't null.
+
+    Reference:
+        https://code.djangoproject.com/ticket/32252
     """
     lhs_mql = process_lhs(self, compiler, connection)
     rhs_mql = process_rhs(self, compiler, connection)
+
     # Get the root column.
     previous = self.lhs
     while isinstance(previous, KeyTransform):
         previous = previous.lhs
     root_column = previous.as_mql(compiler, connection)
-    type_in = _expr_type_in(lhs_mql, ["missing"])
-    result = {"$or": [type_in, {"$eq": [root_column, None]}]}
-    if not rhs_mql:
-        result = {"$not": result}
-    return result
+    return _has_key_predicate(lhs_mql, root_column, negated=rhs_mql)
 
 
 def key_transform_numeric_lookup_mixin(self, compiler, connection):
+    """
+    Checks if the field exists and fulfills the given numeric lookup expression.
+
+    This function ensures that the field in the JSON document:
+    1. Exists (i.e., is not "missing" or "null").
+    2. Satisfies the specified numeric lookup expression.
+
+    Returns:
+        A dictionary representing the MongoDB query that checks both conditions.
+
+    Note:
+        The function constructs the MongoDB expression to check if the field exists
+        and fulfills the numeric lookup expression, and combines these conditions
+        using $and.
+    """
     expr = builtin_lookup(self, compiler, connection)
     lhs = process_lhs(self, compiler, connection)
-    type_in = _expr_type_in(lhs, ["missing", "null"], True)
-    return {"$and": [expr, type_in]}
+    # Check if the type of lhs is not "missing" or "null"
+    not_missing_or_null = {"$not": {"$in": [{"$type": lhs}, ["missing", "null"]]}}
+    return {"$and": [expr, not_missing_or_null]}
 
 
 def register_json_field():