Move Schema key lookup transformer into its own module

django_pydantic_field/v2/fields.py

@@ -5,16 +5,15 @@
 import pydantic
 from django.core import checks, exceptions
 from django.core.serializers.json import DjangoJSONEncoder
-from django.db.models.expressions import BaseExpression, Col, Value
+from django.db.models.expressions import BaseExpression, Value
 from django.db.models.fields import NOT_PROVIDED
 from django.db.models.fields.json import JSONField
-from django.db.models.lookups import Transform
 from django.db.models.query_utils import DeferredAttribute
 
 from django_pydantic_field.compat import deprecation
 from django_pydantic_field.compat.django import GenericContainer
 
-from . import forms, types
+from . import forms, types, lookups
 
 if ty.TYPE_CHECKING:
     import json
@@ -174,7 +173,7 @@ def get_prep_value(self, value: ty.Any):
     def get_transform(self, lookup_name: str):
         transform: ty.Any = super().get_transform(lookup_name)
         if transform is not None:
-            transform = SchemaKeyTransformAdapter(transform)
+            transform = lookups.SchemaKeyTransformAdapter(transform, lookup_name)
         return transform
 
     def get_default(self) -> types.ST:
@@ -198,21 +197,6 @@ def value_to_string(self, obj: Model):
         return self.get_prep_value(value)
 
 
-class SchemaKeyTransformAdapter:
-    """An adapter for creating key transforms for schema field lookups."""
-
-    def __init__(self, transform: type[Transform]):
-        self.transform = transform
-
-    def __call__(self, col: Col | None = None, *args, **kwargs) -> Transform | None:
-        """All transforms should bypass the SchemaField's adaptaion with `get_prep_value`,
-        and routed to JSONField's `get_prep_value` for further processing."""
-        if isinstance(col, BaseExpression):
-            col = col.copy()
-            col.output_field = super(PydanticSchemaField, col.output_field)  # type: ignore
-        return self.transform(col, *args, **kwargs)
-
-
 @ty.overload
 def SchemaField(
     schema: type[types.ST | None] | ty.ForwardRef = ...,

django_pydantic_field/v2/lookups.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import typing as ty
+
+from django.db.models.expressions import BaseExpression, Col
+from django.db.models.lookups import Transform
+
+if ty.TYPE_CHECKING:
+    ExprT = ty.TypeVar("ExprT", bound=BaseExpression)
+
+
+class SchemaKeyTransformAdapter:
+    """An adapter class that modifies the key transformation behavior for schema fields.
+
+    This class acts as an adapter, altering how key transformations are performed on schema fields.
+    It circumvents the usual adaptation process for `PydanticSchemaField` objects,
+    instead opting to use the `JSONField`'s own transformation methods.
+
+    The goal is to utilize the lookup transformation features provided by the `JSONField.encoder` class.
+    With the current limitations in Pydantic, it's not feasible to conduct partial value adaptations.
+
+    While this approach is not ideal for QuerySet lookups,
+    it allows `JSONField.encoder` (which defaults to `DjangoJSONEncoder`) to perform essential transformations.
+    """
+
+    def __init__(self, transform: type[Transform], lookup_name: str):
+        self.transform = transform
+        self.lookup_name = lookup_name
+
+    def __call__(self, col: Col | None = None, *args, **kwargs) -> Transform | None:
+        """All transforms should bypass the SchemaField's adaptaion with `get_prep_value`,
+        and routed to JSONField's `get_prep_value` for further processing."""
+        if isinstance(col, BaseExpression):
+            col = self._get_prep_expression(col)
+        return self.transform(col, *args, **kwargs)
+
+    def _get_prep_expression(self, expr: ExprT) -> ExprT:
+        from .fields import PydanticSchemaField
+
+        if isinstance(expr.output_field, PydanticSchemaField):
+            expr = expr.copy()
+            expr.output_field = super(type(expr.output_field), expr.output_field)  # type: ignore
+        return expr