getField support overlay

Author: Jibola

django_mongodb_backend/query_conversion/expression_converters.py

@@ -10,6 +10,45 @@ class BaseConverter:
     def convert(cls, expr):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @classmethod
+    def is_simple_field_name(cls, field_name):
+        return (
+            isinstance(field_name, str)
+            and field_name != ""
+            and field_name.startswith("$")
+            # Special case for _ separated field names
+            and field_name[1:].replace("_", "").isalnum()
+        )
+
+    @classmethod
+    def is_simple_get_field(cls, get_field_object):
+        if not isinstance(get_field_object, dict):
+            return False
+
+        get_field_expr = get_field_object.get("$getField")
+
+        if (
+            isinstance(get_field_expr, dict)
+            and "input" in get_field_expr
+            and "field" in get_field_expr
+        ):
+            input_expr = get_field_expr["input"]
+            field_name = get_field_expr["field"]
+            return cls.convert_field_name(input_expr) and (
+                isinstance(field_name, str) and not field_name.startswith("$")
+            )
+        return False
+
+    @classmethod
+    def convert_field_name(cls, field_name):
+        if cls.is_simple_field_name(field_name):
+            return field_name[1:]
+        if cls.is_simple_get_field(field_name):
+            get_field_input = field_name["$getField"]["input"]
+            get_field_field = field_name["$getField"]["field"]
+            return f"{cls.convert_field_name(get_field_input)}.{get_field_field}"
+        return None
+
     @classmethod
     def is_simple_value(cls, value):
         """Is the value is a simple type (not a dict)?"""
@@ -19,7 +58,6 @@ def is_simple_value(cls, value):
             return False
         if isinstance(value, (list, tuple, set)):
             return all(cls.is_simple_value(v) for v in value)
-        # TODO: Support `$getField` conversion.
         return not isinstance(value, dict)
 
 
@@ -35,12 +73,7 @@ def convert(cls, args):
         if isinstance(args, list) and len(args) == 2:
             field_expr, value = args
             # Check if first argument is a simple field reference.
-            if (
-                isinstance(field_expr, str)
-                and field_expr.startswith("$")
-                and cls.is_simple_value(value)
-            ):
-                field_name = field_expr[1:]  # Remove the $ prefix.
+            if (field_name := cls.convert_field_name(field_expr)) and cls.is_simple_value(value):
                 if cls.operator == "$eq":
                     return {field_name: value}
                 return {field_name: {cls.operator: value}}
@@ -135,13 +168,12 @@ def convert(cls, in_args):
             field_expr, values = in_args
 
             # Check if first argument is a simple field reference
-            if isinstance(field_expr, str) and field_expr.startswith("$"):
-                field_name = field_expr[1:]  # Remove the $ prefix
-                if isinstance(values, (list, tuple, set)) and all(
-                    cls.is_simple_value(v) for v in values
-                ):
-                    return {field_name: {"$in": values}}
-
+            # Check if second argument is a list of simple values
+            if (field_name := cls.convert_field_name(field_expr)) and (
+                isinstance(values, list | tuple | set)
+                and all(cls.is_simple_value(v) for v in values)
+            ):
+                return {field_name: {"$in": values}}
         return None