pyoq-sql 1.0.2__py3-none-any.whl

pyoq/query/rows.py

@@ -0,0 +1,156 @@
+"""Comparing several values as one.
+
+A composite key is one key, and asking whether a row is among a set of them
+should read that way. A row value puts columns side by side and compares them
+against tuples of values in one predicate, instead of spelling out an OR of
+ANDs that a reader has to reassemble.
+
+The tuples are checked against the columns by position, so a value of the
+wrong type in the wrong place is refused where it is written.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Generic, TypeVar, TypeVarTuple, cast, overload
+
+from pyoq.query.expressions import Condition, Expression, operand
+from pyoq.query.nodes import (
+    BinaryNode,
+    BinaryOperator,
+    RowValueNode,
+    VariadicNode,
+    VariadicOperator,
+)
+
+if TYPE_CHECKING:
+    from pyoq.query.raw import NodeProvider
+
+Values = TypeVarTuple("Values")
+Value1 = TypeVar("Value1")
+Value2 = TypeVar("Value2")
+Value3 = TypeVar("Value3")
+Value4 = TypeVar("Value4")
+Value5 = TypeVar("Value5")
+
+
+class RowValue(Generic[*Values]):
+    """Columns compared together, as wide as the tuples compared against them."""
+
+    __slots__ = ("_families", "_node")
+
+    def __init__(self, expressions: tuple[Expression[object], ...]) -> None:
+        self._families = tuple(expression.family for expression in expressions)
+        self._node = RowValueNode(tuple(expression.node for expression in expressions))
+
+    def in_(self, *rows: tuple[*Values]) -> Condition:
+        """Whether these columns together match any of these rows."""
+        return self._membership(VariadicOperator.IN, rows)
+
+    def not_in(self, *rows: tuple[*Values]) -> Condition:
+        """Whether these columns together match none of these rows."""
+        return self._membership(VariadicOperator.NOT_IN, rows)
+
+    def eq(self, values: tuple[*Values], /) -> Condition:
+        return self._compared(BinaryOperator.EQUAL, values)
+
+    def ne(self, values: tuple[*Values], /) -> Condition:
+        return self._compared(BinaryOperator.NOT_EQUAL, values)
+
+    def gt(self, values: tuple[*Values], /) -> Condition:
+        return self._compared(BinaryOperator.GREATER_THAN, values)
+
+    def ge(self, values: tuple[*Values], /) -> Condition:
+        return self._compared(BinaryOperator.GREATER_OR_EQUAL, values)
+
+    def lt(self, values: tuple[*Values], /) -> Condition:
+        return self._compared(BinaryOperator.LESS_THAN, values)
+
+    def le(self, values: tuple[*Values], /) -> Condition:
+        return self._compared(BinaryOperator.LESS_OR_EQUAL, values)
+
+    def _membership(
+        self,
+        operator: VariadicOperator,
+        rows: tuple[tuple[object, ...], ...],
+    ) -> Condition:
+        operands = (self._node, *(self._row(row) for row in rows))
+        return Condition(VariadicNode(operator, operands))
+
+    def _compared(
+        self,
+        operator: BinaryOperator,
+        values: tuple[object, ...],
+    ) -> Condition:
+        return Condition(BinaryNode(operator, self._node, self._row(values)))
+
+    def _row(self, values: tuple[object, ...]) -> RowValueNode:
+        if len(values) != len(self._families):
+            message = (
+                f"a row of {len(self._families)} columns takes "
+                f"{len(self._families)} values, not {len(values)}"
+            )
+            raise ValueError(message)
+        return RowValueNode(
+            tuple(
+                operand(value, family)
+                for value, family in zip(values, self._families, strict=True)
+            )
+        )
+
+
+@overload
+def row(
+    first: Expression[Value1],
+    second: Expression[Value2],
+    /,
+) -> RowValue[Value1, Value2]: ...
+
+
+@overload
+def row(
+    first: Expression[Value1],
+    second: Expression[Value2],
+    third: Expression[Value3],
+    /,
+) -> RowValue[Value1, Value2, Value3]: ...
+
+
+@overload
+def row(
+    first: Expression[Value1],
+    second: Expression[Value2],
+    third: Expression[Value3],
+    fourth: Expression[Value4],
+    /,
+) -> RowValue[Value1, Value2, Value3, Value4]: ...
+
+
+@overload
+def row(
+    first: Expression[Value1],
+    second: Expression[Value2],
+    third: Expression[Value3],
+    fourth: Expression[Value4],
+    fifth: Expression[Value5],
+    /,
+) -> RowValue[Value1, Value2, Value3, Value4, Value5]: ...
+
+
+@overload
+def row(*expressions: NodeProvider) -> RowValue[*tuple[object, ...]]: ...
+
+
+def row(*expressions: object) -> RowValue[*tuple[object, ...]]:
+    """Several columns compared as one, the way a composite key is one key."""
+    typed = tuple(_expression(expression) for expression in expressions)
+    return RowValue[*tuple[object, ...]](typed)
+
+
+def _expression(value: object) -> Expression[object]:
+    if isinstance(value, Expression):
+        return cast("Expression[object]", value)
+    message = "a row value is built from typed expressions"
+    raise TypeError(message)
+
+
+__all__ = ("RowValue", "row")