async-orm-kit 1.0.0__py3-none-any.whl

async_orm_kit/query/expressions/collections.py

@@ -0,0 +1,251 @@
+"""Typed JSON and PostgreSQL ARRAY expressions.
+
+These expressions resolve mapped columns through the normal field resolver; JSON
+path selectors use SQLAlchemy's dialect-aware JSON operators (not string-built
+SQL). PostgreSQL-only operators require PostgreSQL column types and compile with
+bound parameters. Unsupported combinations raise instead of falling back to
+LIKE comparisons or incorrect indexes.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from sqlalchemy import Integer, func
+from sqlalchemy.dialects.postgresql import ARRAY, JSONB
+from sqlalchemy.exc import CompileError
+from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.sql.elements import ColumnElement
+from sqlalchemy.sql.sqltypes import JSON
+from sqlalchemy.sql.visitors import InternalTraversal
+
+from async_orm_kit.query.expressions.functions import _operand
+from async_orm_kit.query.lookups.resolver import ResolverContext
+from async_orm_kit.query.mixins.combinable import CombinableMixin
+from async_orm_kit.query.mixins.orderable import OrderableMixin
+
+
+class _PostgreSQLOnly(ColumnElement):
+    """Prevent PostgreSQL-only predicates from silently compiling elsewhere."""
+
+    inherit_cache = True
+    _traverse_internals = [("_wrapped", InternalTraversal.dp_clauseelement)]
+
+    def __init__(self, expression):
+        self._wrapped = expression
+        self.type = expression.type
+
+
+@compiles(_PostgreSQLOnly, "postgresql")
+def _compile_postgresql(element, compiler, **kwargs):
+    return compiler.process(element._wrapped, **kwargs)
+
+
+@compiles(_PostgreSQLOnly)
+def _reject_other_dialects(element, compiler, **kwargs):
+    raise CompileError("This JSONB/ARRAY operation is supported only on PostgreSQL")
+
+
+class _TypedExpression(CombinableMixin, OrderableMixin):
+    __slots__ = ()
+
+
+def _json_source(expression: Any, resolver: Any, context: ResolverContext):
+    source = _operand(expression, resolver, context)
+    if not isinstance(source.type, JSON):
+        raise TypeError("JSON expressions require a SQLAlchemy JSON/JSONB column")
+    return source
+
+
+def _postgres_array(expression: Any, resolver: Any, context: ResolverContext):
+    source = _operand(expression, resolver, context)
+    if not isinstance(source.type, ARRAY):
+        raise TypeError("ARRAY expressions require a PostgreSQL ARRAY column")
+    return source
+
+
+class JSONPath(_TypedExpression):
+    """Extract a JSON path, retaining an explicit result type.
+
+    ``JSONPath(F('payload'), 'customer', 'id', as_type='integer')`` compiles
+    with the JSON operators of the active SQLAlchemy dialect. The path is a
+    SQLAlchemy JSON index parameter, not interpolated SQL text.
+
+    ``as_type`` can be ``json``, ``string``, ``integer``, ``float`` or
+    ``boolean``. JSON null and SQL NULL have distinct semantics and should be
+    tested for each intended dialect.
+    """
+
+    __slots__ = ("expression", "path", "as_type")
+    _types = frozenset({"json", "string", "integer", "float", "boolean"})
+
+    def __init__(
+        self,
+        expression: Any,
+        *path: str | int,
+        as_type: Literal["json", "string", "integer", "float", "boolean"] = "json",
+    ):
+        if not path:
+            raise ValueError("JSONPath requires at least one path segment")
+        if any(
+            isinstance(part, bool)
+            or not isinstance(part, (str, int))
+            or (isinstance(part, str) and not part)
+            or (isinstance(part, int) and part < 0)
+            for part in path
+        ):
+            raise ValueError(
+                "JSONPath segments must be non-empty keys or nonnegative indexes"
+            )
+        if as_type not in self._types:
+            raise ValueError(f"Unsupported JSON result type: {as_type!r}")
+        self.expression = expression
+        self.path = tuple(path)
+        self.as_type = as_type
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        source = _json_source(self.expression, field_resolver, context)
+        # Multi-segment JSON paths compile to #> for PostgreSQL and
+        # json_extract/JSON_EXTRACT for SQLite/MySQL.
+        selected = source[self.path[0] if len(self.path) == 1 else self.path]
+        if self.as_type == "json":
+            return selected
+        return getattr(selected, f"as_{self.as_type}")()
+
+
+class JSONHasKey(_TypedExpression):
+    """PostgreSQL JSONB top-level key-existence operator ``?``.
+
+    The JSONB comparator is required to prevent SQLAlchemy's generic string
+    containment behavior from being used on SQLite or MySQL JSON values.
+    """
+
+    __slots__ = ("expression", "key")
+
+    def __init__(self, expression: Any, key: str):
+        if not isinstance(key, str) or not key:
+            raise ValueError("JSONHasKey requires a non-empty key")
+        self.expression = expression
+        self.key = key
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        source = _json_source(self.expression, field_resolver, context)
+        if not isinstance(source.type, JSONB):
+            raise TypeError("JSONHasKey requires a PostgreSQL JSONB column")
+        return _PostgreSQLOnly(source.has_key(self.key))  # noqa: W601
+
+
+class JSONContains(_TypedExpression):
+    """PostgreSQL JSONB containment ``@>`` with a bound JSON value."""
+
+    __slots__ = ("expression", "value")
+
+    def __init__(self, expression: Any, value: dict | list):
+        if not isinstance(value, (dict, list)):
+            raise TypeError("JSONContains expects a JSON object or array")
+        self.expression = expression
+        self.value = value
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        source = _json_source(self.expression, field_resolver, context)
+        if not isinstance(source.type, JSONB):
+            raise TypeError("JSONContains requires a PostgreSQL JSONB column")
+        return _PostgreSQLOnly(source.contains(self.value))
+
+
+class ArrayIndex(_TypedExpression):
+    """Index a PostgreSQL ARRAY using the SQLAlchemy column's index convention.
+
+    Normally PostgreSQL arrays are 1-indexed. If the mapped SQLAlchemy ARRAY
+    has ``zero_indexes=True``, SQLAlchemy adjusts indexes automatically.
+    """
+
+    __slots__ = ("expression", "index")
+
+    def __init__(self, expression: Any, index: int):
+        if isinstance(index, bool) or not isinstance(index, int) or index < 0:
+            raise ValueError("ArrayIndex requires a nonnegative integer")
+        self.expression = expression
+        self.index = index
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        source = _postgres_array(self.expression, field_resolver, context)
+        if self.index == 0 and not source.type.zero_indexes:
+            raise ValueError(
+                "PostgreSQL ARRAY indexes start at 1 unless zero_indexes=True"
+            )
+        return _PostgreSQLOnly(source[self.index])
+
+
+class ArrayContains(_TypedExpression):
+    """PostgreSQL ARRAY containment ``@>``."""
+
+    __slots__ = ("expression", "values")
+
+    def __init__(self, expression: Any, values: list | tuple):
+        if not isinstance(values, (list, tuple)):
+            raise TypeError("ArrayContains expects a list or tuple")
+        self.expression = expression
+        self.values = list(values)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return _PostgreSQLOnly(
+            _postgres_array(self.expression, field_resolver, context).contains(
+                self.values
+            )
+        )
+
+
+class ArrayOverlap(_TypedExpression):
+    """PostgreSQL ARRAY overlap ``&&``."""
+
+    __slots__ = ("expression", "values")
+
+    def __init__(self, expression: Any, values: list | tuple):
+        if not isinstance(values, (list, tuple)):
+            raise TypeError("ArrayOverlap expects a list or tuple")
+        self.expression = expression
+        self.values = list(values)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return _PostgreSQLOnly(
+            _postgres_array(self.expression, field_resolver, context).overlap(
+                self.values
+            )
+        )
+
+
+class ArrayLength(_TypedExpression):
+    """PostgreSQL ``array_length(array, dimension)``."""
+
+    __slots__ = ("expression", "dimension")
+
+    def __init__(self, expression: Any, dimension: int = 1):
+        if (
+            isinstance(dimension, bool)
+            or not isinstance(dimension, int)
+            or dimension < 1
+        ):
+            raise ValueError("ArrayLength dimension must be a positive integer")
+        self.expression = expression
+        self.dimension = dimension
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        return _PostgreSQLOnly(
+            func.array_length(
+                _postgres_array(self.expression, field_resolver, context),
+                self.dimension,
+                type_=Integer(),
+            )
+        )
+
+
+__all__ = [
+    "ArrayContains",
+    "ArrayIndex",
+    "ArrayLength",
+    "ArrayOverlap",
+    "JSONContains",
+    "JSONHasKey",
+    "JSONPath",
+]

async_orm_kit/query/expressions/f.py

@@ -0,0 +1,362 @@
+import operator
+from typing import Any
+
+from sqlalchemy import func
+from sqlalchemy.sql.elements import ClauseElement, literal
+
+from async_orm_kit.query.expressions.case import Case
+from async_orm_kit.query.expressions.value import Value
+from async_orm_kit.query.lookups.resolver import ResolverContext
+from async_orm_kit.query.mixins.combinable import CombinableMixin
+from async_orm_kit.query.mixins.orderable import OrderableMixin
+from async_orm_kit.types import SAFE_SCALAR_TYPES
+
+
+class F(CombinableMixin, OrderableMixin):
+    """Reference to a model field for database-level operations and expressions.
+
+    F objects enable database-side computation by referencing field names that get
+    resolved to actual columns at query time. They support arithmetic operations,
+    comparisons, and can traverse relationships using double-underscore syntax.
+    Operations on F objects create FExpression instances that resolve to SQLAlchemy
+    expressions, allowing efficient database-side calculations without loading data
+    into Python.
+
+    Common use cases include updating fields based on their current values, computing
+    derived fields in annotations, and building complex filter conditions that compare
+    fields to each other rather than to literal values.
+
+    Args:
+        name: Field name or relationship path using double-underscore syntax.
+              Examples: "price", "quantity", "order__customer__name"
+
+    Note:
+        F objects are lazy - they don't execute queries until the QuerySet is evaluated.
+        Operations create expression trees that resolve to SQL at query compilation time.
+        Relationship traversal automatically registers joins with the QueryBuilder,
+        ensuring proper JOIN clauses in the generated SQL.
+    """
+
+    _name: str
+    __slots__ = ("_name",)
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        raise AttributeError("F expressions are immutable.")
+
+    def __init__(self, name: str) -> None:
+        """Initialize F object with field name or relationship path.
+
+        Args:
+            name: Field name or relationship path.
+             Use double-underscore for relationships: "field_name" or "relation__field_name"
+
+        Example:
+            >>> F("price")
+            >>> F("order__customer__email")
+        """
+        if not isinstance(name, str) or not name:
+            raise ValueError("F() requires a non-empty field name.")
+        object.__setattr__(self, "_name", name)
+
+    def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
+        """Resolve field name to SQLAlchemy column expression.
+
+        Uses FieldResolver to handle relationship traversal and join registration.
+        The resolver maintains join state across the entire query, ensuring F objects
+        in different parts of the query (filters, annotations, order_by) share the
+        same join strategy.
+
+        Args:
+            field_resolver: FieldResolver instance for column resolution and join tracking.
+            context: ResolverContext for OR context and filter marking propagation.
+
+        Returns:
+            SQLAlchemy Column object representing the resolved field.
+
+        Example:
+            >>> # Internal usage during query compilation
+            >>> f = F("order__customer__name")
+            >>> column = f.resolve(field_resolver, context)
+            >>> # Returns Customer.name column with joins registered
+        """
+        resolution = field_resolver.resolve_to_column(key=self.name, context=context)
+        return resolution.column
+
+    def __pow__(self, other):
+        """Power operation: F("x") ** 2 creates exponentiation expression."""
+        return FExpression(self, "pow", other)
+
+    def __rpow__(self, other):
+        """Reverse power: 2 ** F("x") creates exponentiation expression."""
+        return FExpression(other, "pow", self)
+
+    def __eq__(self, other):
+        """Equality comparison: F("x") == value creates comparison expression."""
+        return FExpression(self, operator.eq, other)
+
+    def __ne__(self, other):
+        """Inequality comparison: F("x") != value creates comparison expression."""
+        return FExpression(self, operator.ne, other)
+
+    def __lt__(self, other):
+        """Less than: F("x") < value creates comparison expression."""
+        return FExpression(self, operator.lt, other)
+
+    def __le__(self, other):
+        """Less than or equal: F("x") <= value creates comparison expression."""
+        return FExpression(self, operator.le, other)
+
+    def __gt__(self, other):
+        """Greater than: F("x") > value creates comparison expression."""
+        return FExpression(self, operator.gt, other)
+
+    def __ge__(self, other):
+        """Greater than or equal: F("x") >= value creates comparison expression."""
+        return FExpression(self, operator.ge, other)
+
+    def __pos__(self):
+        """Unary plus: +F("x") returns self unchanged."""
+        return self
+
+    def __repr__(self) -> str:
+        """Return string representation for debugging."""
+        return f"F('{self.name}')"
+
+    def __hash__(self) -> int:
+        """Return hash based on field name for use in sets and dicts."""
+        return hash(self.name)
+
+
+class FExpression(CombinableMixin, OrderableMixin):
+    """Complex expression tree built from F objects and operations.
+
+    FExpression represents the result of operations on F objects, creating an
+    expression tree that resolves to SQLAlchemy expressions at query time. Supports
+    arithmetic operations (``+``, ``-``, ``*``, ``/``, ``//``, ``%``, ``**``),
+    comparisons (``==``, ``!=``, ``<``, ``<=``, ``>``, ``>=``),
+    bitwise operations (``&``, ``|``, ``^``), and unary operations (-, ~, abs).
+
+    Expressions can be nested arbitrarily deep and combined with literals, other F
+    objects, Value() wrappers, Case() conditionals, and Aggregate functions. The
+    expression tree is evaluated lazily during query compilation, generating efficient
+    SQL that performs all computation database-side.
+
+    Args:
+        left: Left operand - can be F, FExpression, Value, Case, Aggregate, scalar,
+              or None for unary operations.
+        op: Operator function from operator module or special string like "pow", "abs".
+        right: Right operand - same types as left.
+
+    Note:
+        Expressions are immutable - each operation creates a new FExpression instance.
+        This enables safe reuse of expression components across multiple queries.
+
+        Resolution is recursive - complex nested expressions resolve from leaves to
+        root, with each operand resolving independently before the operation is applied.
+    """
+
+    _left: Any
+    _operator: Any
+    _right: Any
+    __slots__ = ("_left", "_operator", "_right")
+
+    @property
+    def left(self):
+        return self._left
+
+    @property
+    def operator(self):
+        return self._operator
+
+    @property
+    def right(self):
+        return self._right
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        raise AttributeError("FExpression objects are immutable.")
+
+    def __init__(self, left, op, right):
+        """Initialize expression with operands and operator.
+
+        Args:
+            left: Left operand or None for unary operations.
+            op: Operator function or special string ("pow", "abs").
+            right: Right operand.
+        """
+        allowed_operators = {
+            operator.add,
+            operator.sub,
+            operator.mul,
+            operator.truediv,
+            operator.mod,
+            operator.neg,
+            operator.eq,
+            operator.ne,
+            operator.lt,
+            operator.le,
+            operator.gt,
+            operator.ge,
+            "pow",
+            "abs",
+            "floordiv",
+            "bit_and",
+            "bit_or",
+            "bit_not",
+        }
+        if op not in allowed_operators:
+            raise ValueError(f"Unsupported F expression operator: {op!r}")
+        object.__setattr__(self, "_left", left)
+        object.__setattr__(self, "_operator", op)
+        object.__setattr__(self, "_right", right)
+
+    def resolve(
+        self,
+        field_resolver,
+        context: ResolverContext = ResolverContext(),
+    ):
+        """Resolve expression tree to SQLAlchemy expression.
+
+        Recursively resolves operands then applies the operator.
+        Handles special function-based operations (abs, pow), unary operations (left=None), and
+        standard binary operations.
+
+        Args:
+            field_resolver: FieldResolver instance for column resolution and join tracking.
+            context: ResolverContext for OR context and filter marking propagation.
+
+        Returns:
+            SQLAlchemy expression representing the resolved operation.
+
+        Example:
+            >>> # Internal usage during query compilation
+            >>> expr = F("price") * F("quantity")
+            >>> sql_expr = expr.resolve(field_resolver, context)
+            >>> # Returns: products.price * products.quantity
+        """
+        left_val = self._resolve_operand(self.left, field_resolver, context)
+        right_val = self._resolve_operand(self.right, field_resolver, context)
+
+        # Handle special function-based operations
+        if self.operator == "abs":
+            return func.abs(right_val)
+        if self.operator == "floordiv":
+            return func.floor(left_val / right_val)
+        if self.operator == "bit_and":
+            return left_val.bitwise_and(right_val)
+        if self.operator == "bit_or":
+            return left_val.bitwise_or(right_val)
+        if self.operator == "bit_not":
+            return right_val.bitwise_not()
+        elif self.operator == "pow":
+            return func.power(left_val, right_val)
+
+        # Handle unary operations
+        if self.left is None:
+            return self.operator(right_val)
+
+        # Handle binary operations
+        return self.operator(left_val, right_val)
+
+    def _resolve_operand(
+        self, operand, field_resolver, context: ResolverContext = ResolverContext()
+    ):
+        """Resolve operand to SQLAlchemy expression or literal.
+
+        Handles multiple operand types with priority-based resolution:
+        1. None - Returns None for unary operations
+        2. F/FExpression - Resolves to column/expression
+        3. Value/Case - Resolves to literal/conditional
+        4. Aggregate - Resolves to aggregate function
+        5. ClauseElement - SQLAlchemy expression passed through
+        6. Safe scalars - Wrapped in literal()
+        7. Other types - Raises TypeError
+
+        Args:
+            operand: Value to resolve.
+            field_resolver: FieldResolver for column resolution.
+            context: ResolverContext for propagation.
+
+        Returns:
+            SQLAlchemy expression or literal value.
+
+        Raises:
+            TypeError: If operand type is unsupported.
+        """
+        from async_orm_kit.query.aggregations.base import Aggregate
+
+        if operand is None:
+            return None
+
+        if isinstance(operand, (F, FExpression)):
+            return operand.resolve(field_resolver, context)
+
+        if isinstance(operand, (Value, Case)):
+            return operand.resolve(field_resolver, context)
+
+        if isinstance(operand, Aggregate):
+            return operand.resolve(field_resolver=field_resolver, context=context)
+
+        if isinstance(operand, CombinableMixin) and hasattr(operand, "resolve"):
+            return operand.resolve(field_resolver=field_resolver, context=context)
+
+        if isinstance(operand, ClauseElement):
+            return operand
+
+        if isinstance(operand, SAFE_SCALAR_TYPES):
+            return literal(operand)
+
+        raise TypeError(
+            f"Unsupported operand type in F expression: {type(operand).__name__!r}. "
+            f"Use F(), Value(), Case(), or a scalar (int, float, str, etc.)."
+        )
+
+    # ── Operations NOT provided by CombinableMixin ──
+
+    def __pow__(self, other):
+        """Power operation: expr ** 2 creates nested exponentiation."""
+        return FExpression(self, "pow", other)
+
+    def __rpow__(self, other):
+        """Reverse power: 2 ** expr creates nested exponentiation."""
+        return FExpression(other, "pow", self)
+
+    def __eq__(self, other):
+        """Equality comparison: expr == value creates comparison."""
+        return FExpression(self, operator.eq, other)
+
+    def __ne__(self, other):
+        """Inequality comparison: expr != value creates comparison."""
+        return FExpression(self, operator.ne, other)
+
+    def __lt__(self, other):
+        """Less than: expr < value creates comparison."""
+        return FExpression(self, operator.lt, other)
+
+    def __le__(self, other):
+        """Less than or equal: expr <= value creates comparison."""
+        return FExpression(self, operator.le, other)
+
+    def __gt__(self, other):
+        """Greater than: expr > value creates comparison."""
+        return FExpression(self, operator.gt, other)
+
+    def __ge__(self, other):
+        """Greater than or equal: expr >= value creates comparison."""
+        return FExpression(self, operator.ge, other)
+
+    def __hash__(self) -> int:
+        """Return hash based on operand identities for use in sets and dicts."""
+        return hash((id(self.left), id(self.operator), id(self.right)))
+
+    def __repr__(self) -> str:
+        """Return string representation showing expression tree structure."""
+        op_name = (
+            self.operator.__name__
+            if hasattr(self.operator, "__name__")
+            else str(self.operator)
+        )
+        return f"FExpression({self.left!r}, {op_name}, {self.right!r})"