From 781ac8f0aa1ef9289e424f451353f491b09bd01f Mon Sep 17 00:00:00 2001 From: jazzthief Date: Thu, 2 Feb 2023 13:48:13 -0500 Subject: [PATCH] Dedicated bitwise operators Added a full suite of new SQL bitwise operators, for performing database-side bitwise expressions on appropriate data values such as integers, bit-strings, and similar. Pull request courtesy Yegor Statkevich. Fixes: #8780 Closes: #9204 Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/9204 Pull-request-sha: a4541772a6a784f9161ad78ef84d2ea7a62fa8de Change-Id: I4c70e80f9548dcc1b4e3dccd71bd59d51d3ed46e --- doc/build/changelog/unreleased_20/8780.rst | 12 ++ doc/build/core/operators.rst | 67 +++++++++++ doc/build/core/sqlelement.rst | 2 + lib/sqlalchemy/__init__.py | 1 + lib/sqlalchemy/dialects/postgresql/base.py | 3 + lib/sqlalchemy/sql/_elements_constructors.py | 17 +++ lib/sqlalchemy/sql/compiler.py | 7 ++ lib/sqlalchemy/sql/default_comparator.py | 16 +++ lib/sqlalchemy/sql/elements.py | 15 +++ lib/sqlalchemy/sql/expression.py | 1 + lib/sqlalchemy/sql/operators.py | 120 +++++++++++++++++++ test/dialect/postgresql/test_compiler.py | 8 ++ test/sql/test_operators.py | 71 +++++++++++ 13 files changed, 340 insertions(+) create mode 100644 doc/build/changelog/unreleased_20/8780.rst diff --git a/doc/build/changelog/unreleased_20/8780.rst b/doc/build/changelog/unreleased_20/8780.rst new file mode 100644 index 0000000000..22c20c2b2a --- /dev/null +++ b/doc/build/changelog/unreleased_20/8780.rst @@ -0,0 +1,12 @@ +.. change:: + :tags: usecase, sql + :tickets: 8780 + + Added a full suite of new SQL bitwise operators, for performing + database-side bitwise expressions on appropriate data values such as + integers, bit-strings, and similar. Pull request courtesy Yegor Statkevich. + + .. seealso:: + + :ref:`operators_bitwise` + diff --git a/doc/build/core/operators.rst b/doc/build/core/operators.rst index f75413b3c2..0450aab03e 100644 --- a/doc/build/core/operators.rst +++ b/doc/build/core/operators.rst @@ -583,6 +583,73 @@ Arithmetic Operators .. +.. _operators_bitwise: + +Bitwise Operators +^^^^^^^^^^^^^^^^^ + +Bitwise operator functions provide uniform access to bitwise operators across +different backends, which are expected to operate on compatible +values such as integers and bit-strings (e.g. PostgreSQL +:class:`_postgresql.BIT` and similar). Note that these are **not** general +boolean operators. + +.. versionadded:: 2.0.2 Added dedicated operators for bitwise operations. + +* :meth:`_sql.ColumnOperators.bitwise_not`, :func:`_sql.bitwise_not`. + Available as a column-level method, producing a bitwise NOT clause against a + parent object:: + + >>> print(column("x").bitwise_not()) + ~x + + This operator is also available as a column-expression-level method, applying + bitwise NOT to an individual column expression:: + + >>> from sqlalchemy import bitwise_not + >>> print(bitwise_not(column("x"))) + ~x + + .. + +* :meth:`_sql.ColumnOperators.bitwise_and` produces bitwise AND:: + + >>> print(column("x").bitwise_and(5)) + x & :x_1 + + .. + +* :meth:`_sql.ColumnOperators.bitwise_or` produces bitwise OR:: + + >>> print(column("x").bitwise_or(5)) + x | :x_1 + + .. + +* :meth:`_sql.ColumnOperators.bitwise_xor` produces bitwise XOR:: + + >>> print(column("x").bitwise_xor(5)) + x ^ :x_1 + + For PostgreSQL dialects, "#" is used to represent bitwise XOR; this emits + automatically when using one of these backends:: + + >>> from sqlalchemy.dialects import postgresql + >>> print(column("x").bitwise_xor(5).compile(dialect=postgresql.dialect())) + x # %(x_1)s + + .. + +* :meth:`_sql.ColumnOperators.bitwise_rshift`, :meth:`_sql.ColumnOperators.bitwise_lshift` + produce bitwise shift operators:: + + >>> print(column("x").bitwise_rshift(5)) + x >> :x_1 + >>> print(column("x").bitwise_lshift(5)) + x << :x_1 + + .. + Using Conjunctions and Negations ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/doc/build/core/sqlelement.rst b/doc/build/core/sqlelement.rst index be780adb0e..ac55e988ed 100644 --- a/doc/build/core/sqlelement.rst +++ b/doc/build/core/sqlelement.rst @@ -26,6 +26,8 @@ used when building up SQLAlchemy Expression Language constructs. .. autofunction:: bindparam +.. autofunction:: bitwise_not + .. autofunction:: case .. autofunction:: cast diff --git a/lib/sqlalchemy/__init__.py b/lib/sqlalchemy/__init__.py index 447451f4b4..9cb8140c6c 100644 --- a/lib/sqlalchemy/__init__.py +++ b/lib/sqlalchemy/__init__.py @@ -89,6 +89,7 @@ from .sql.expression import between as between from .sql.expression import BinaryExpression as BinaryExpression from .sql.expression import bindparam as bindparam from .sql.expression import BindParameter as BindParameter +from .sql.expression import bitwise_not as bitwise_not from .sql.expression import BooleanClauseList as BooleanClauseList from .sql.expression import CacheKey as CacheKey from .sql.expression import Case as Case diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py index d7e6634b31..255c72042e 100644 --- a/lib/sqlalchemy/dialects/postgresql/base.py +++ b/lib/sqlalchemy/dialects/postgresql/base.py @@ -1730,6 +1730,9 @@ class PGCompiler(compiler.SQLCompiler): self.process(element.stop, **kw), ) + def visit_bitwise_xor_op_binary(self, binary, operator, **kw): + return self._generate_generic_binary(binary, " # ", **kw) + def visit_json_getitem_op_binary( self, binary, operator, _cast_applied=False, **kw ): diff --git a/lib/sqlalchemy/sql/_elements_constructors.py b/lib/sqlalchemy/sql/_elements_constructors.py index 9b96322734..98f5a1cc6f 100644 --- a/lib/sqlalchemy/sql/_elements_constructors.py +++ b/lib/sqlalchemy/sql/_elements_constructors.py @@ -1030,6 +1030,23 @@ def distinct(expr: _ColumnExpressionArgument[_T]) -> UnaryExpression[_T]: return UnaryExpression._create_distinct(expr) +def bitwise_not(expr: _ColumnExpressionArgument[_T]) -> UnaryExpression[_T]: + """Produce a unary bitwise NOT clause, typically via the ``~`` operator. + + Not to be confused with boolean negation :func:`_sql.not_`. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + + """ + + return UnaryExpression._create_bitwise_not(expr) + + def extract(field: str, expr: _ColumnExpressionArgument[Any]) -> Extract: """Return a :class:`.Extract` construct. diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py index d4ddc2e5da..bc463f9a1e 100644 --- a/lib/sqlalchemy/sql/compiler.py +++ b/lib/sqlalchemy/sql/compiler.py @@ -279,6 +279,13 @@ OPERATORS = { operators.asc_op: " ASC", operators.nulls_first_op: " NULLS FIRST", operators.nulls_last_op: " NULLS LAST", + # bitwise + operators.bitwise_xor_op: " ^ ", + operators.bitwise_or_op: " | ", + operators.bitwise_and_op: " & ", + operators.bitwise_not_op: "~", + operators.bitwise_lshift_op: " << ", + operators.bitwise_rshift_op: " >> ", } FUNCTIONS: Dict[Type[Function[Any]], str] = { diff --git a/lib/sqlalchemy/sql/default_comparator.py b/lib/sqlalchemy/sql/default_comparator.py index 19b5291c4f..57460b0365 100644 --- a/lib/sqlalchemy/sql/default_comparator.py +++ b/lib/sqlalchemy/sql/default_comparator.py @@ -273,6 +273,16 @@ def _neg_impl( return UnaryExpression(expr, operator=operators.neg, type_=expr.type) +def _bitwise_not_impl( + expr: ColumnElement[Any], op: OperatorType, **kw: Any +) -> ColumnElement[Any]: + """See :meth:`.ColumnOperators.bitwise_not`.""" + + return UnaryExpression( + expr, operator=operators.bitwise_not_op, type_=expr.type + ) + + def _match_impl( expr: ColumnElement[Any], op: OperatorType, other: Any, **kw: Any ) -> ColumnElement[Any]: @@ -420,6 +430,12 @@ operator_lookup: Dict[ "sub": (_binary_operate, util.EMPTY_DICT), "div": (_binary_operate, util.EMPTY_DICT), "mod": (_binary_operate, util.EMPTY_DICT), + "bitwise_xor_op": (_binary_operate, util.EMPTY_DICT), + "bitwise_or_op": (_binary_operate, util.EMPTY_DICT), + "bitwise_and_op": (_binary_operate, util.EMPTY_DICT), + "bitwise_not_op": (_bitwise_not_impl, util.EMPTY_DICT), + "bitwise_lshift_op": (_binary_operate, util.EMPTY_DICT), + "bitwise_rshift_op": (_binary_operate, util.EMPTY_DICT), "truediv": (_binary_operate, util.EMPTY_DICT), "floordiv": (_binary_operate, util.EMPTY_DICT), "custom_op": (_custom_op_operate, util.EMPTY_DICT), diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py index c3f7b884ce..70c65b5a1a 100644 --- a/lib/sqlalchemy/sql/elements.py +++ b/lib/sqlalchemy/sql/elements.py @@ -3587,6 +3587,21 @@ class UnaryExpression(ColumnElement[_T]): wraps_column_expression=False, ) + @classmethod + def _create_bitwise_not( + cls, + expr: _ColumnExpressionArgument[_T], + ) -> UnaryExpression[_T]: + col_expr: ColumnElement[_T] = coercions.expect( + roles.ExpressionElementRole, expr + ) + return UnaryExpression( + col_expr, + operator=operators.bitwise_not_op, + type_=col_expr.type, + wraps_column_expression=False, + ) + @property def _order_by_label_element(self) -> Optional[Label[Any]]: if self.modifier in (operators.desc_op, operators.asc_op): diff --git a/lib/sqlalchemy/sql/expression.py b/lib/sqlalchemy/sql/expression.py index 3f5c76a507..7076cd10d2 100644 --- a/lib/sqlalchemy/sql/expression.py +++ b/lib/sqlalchemy/sql/expression.py @@ -22,6 +22,7 @@ from ._elements_constructors import any_ as any_ from ._elements_constructors import asc as asc from ._elements_constructors import between as between from ._elements_constructors import bindparam as bindparam +from ._elements_constructors import bitwise_not as bitwise_not from ._elements_constructors import case as case from ._elements_constructors import cast as cast from ._elements_constructors import collate as collate diff --git a/lib/sqlalchemy/sql/operators.py b/lib/sqlalchemy/sql/operators.py index 567802916f..c973126ca4 100644 --- a/lib/sqlalchemy/sql/operators.py +++ b/lib/sqlalchemy/sql/operators.py @@ -697,6 +697,90 @@ class ColumnOperators(Operators): """ return self.operate(ilike_op, other, escape=escape) + def bitwise_xor(self, other: Any) -> ColumnOperators: + """Produce a bitwise XOR operation, typically via the ``^`` + operator, or ``#`` for PostgreSQL. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + """ + + return self.operate(bitwise_xor_op, other) + + def bitwise_or(self, other: Any) -> ColumnOperators: + """Produce a bitwise OR operation, typically via the ``|`` + operator. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + """ + + return self.operate(bitwise_or_op, other) + + def bitwise_and(self, other: Any) -> ColumnOperators: + """Produce a bitwise AND operation, typically via the ``&`` + operator. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + """ + + return self.operate(bitwise_and_op, other) + + def bitwise_not(self) -> ColumnOperators: + """Produce a bitwise NOT operation, typically via the ``~`` + operator. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + """ + + return self.operate(bitwise_not_op) + + def bitwise_lshift(self, other: Any) -> ColumnOperators: + """Produce a bitwise LSHIFT operation, typically via the ``<<`` + operator. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + """ + + return self.operate(bitwise_lshift_op, other) + + def bitwise_rshift(self, other: Any) -> ColumnOperators: + """Produce a bitwise RSHIFT operation, typically via the ``>>`` + operator. + + .. versionadded:: 2.0.2 + + .. seealso:: + + :ref:`operators_bitwise` + + """ + + return self.operate(bitwise_rshift_op, other) + def in_(self, other: Any) -> ColumnOperators: """Implement the ``in`` operator. @@ -2266,6 +2350,36 @@ def json_path_getitem_op(a: Any, b: Any) -> Any: raise NotImplementedError() +@_operator_fn +def bitwise_xor_op(a: Any, b: Any) -> Any: + return a.bitwise_xor(b) + + +@_operator_fn +def bitwise_or_op(a: Any, b: Any) -> Any: + return a.bitwise_or(b) + + +@_operator_fn +def bitwise_and_op(a: Any, b: Any) -> Any: + return a.bitwise_and(b) + + +@_operator_fn +def bitwise_not_op(a: Any) -> Any: + return a.bitwise_not() + + +@_operator_fn +def bitwise_lshift_op(a: Any, b: Any) -> Any: + return a.bitwise_lshift(b) + + +@_operator_fn +def bitwise_rshift_op(a: Any, b: Any) -> Any: + return a.bitwise_rshift(b) + + def is_comparison(op: OperatorType) -> bool: return op in _comparison or isinstance(op, custom_op) and op.is_comparison @@ -2344,8 +2458,14 @@ _PRECEDENCE: Dict[OperatorType, int] = { floordiv: 8, mod: 8, neg: 8, + bitwise_not_op: 8, add: 7, sub: 7, + bitwise_xor_op: 7, + bitwise_or_op: 7, + bitwise_and_op: 7, + bitwise_lshift_op: 7, + bitwise_rshift_op: 7, concat_op: 6, filter_op: 6, match_op: 5, diff --git a/test/dialect/postgresql/test_compiler.py b/test/dialect/postgresql/test_compiler.py index 57b147c90d..080cfb767d 100644 --- a/test/dialect/postgresql/test_compiler.py +++ b/test/dialect/postgresql/test_compiler.py @@ -2424,6 +2424,14 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL): "WHERE usages.date <@ %(date_1)s::DATERANGE", ) + def test_bitwise_xor(self): + c1 = column("c1", Integer) + c2 = column("c2", Integer) + self.assert_compile( + select(c1.bitwise_xor(c2)), + "SELECT c1 # c2 AS anon_1", + ) + class InsertOnConflictTest(fixtures.TablesTest, AssertsCompiledSQL): __dialect__ = postgresql.dialect() diff --git a/test/sql/test_operators.py b/test/sql/test_operators.py index d93ba61bac..fd0cf66549 100644 --- a/test/sql/test_operators.py +++ b/test/sql/test_operators.py @@ -1783,6 +1783,35 @@ class OperatorPrecedenceTest(fixtures.TestBase, testing.AssertsCompiledSQL): self.assert_compile(op2, "mytable.myid hoho :myid_1 lala :param_1") self.assert_compile(op3, "(mytable.myid hoho :myid_1) lala :param_1") + def test_bitwise_not_precedence(self): + op1 = operators.bitwise_not_op + c = self.table1.c.myid + op2 = op1(c).op("lala", precedence=7)(4) + op3 = op1(c).op("lala", precedence=9)(4) + + self.assert_compile(op2, "~mytable.myid lala :param_1") + self.assert_compile(op3, "(~mytable.myid) lala :param_1") + + @testing.combinations( + ("xor", operators.bitwise_xor_op, "^"), + ("or", operators.bitwise_or_op, "|"), + ("and", operators.bitwise_and_op, "&"), + ("lshift", operators.bitwise_lshift_op, "<<"), + ("rshift", operators.bitwise_rshift_op, ">>"), + id_="iaa", + ) + def test_bitwise_op_precedence(self, py_op, sql_op): + c = self.table1.c.myid + op1 = py_op(c, 5).op("lala", precedence=6)(4) + op2 = py_op(c, 5).op("lala", precedence=8)(4) + + self.assert_compile( + op1, f"mytable.myid {sql_op} :myid_1 lala :param_1" + ) + self.assert_compile( + op2, f"(mytable.myid {sql_op} :myid_1) lala :param_1" + ) + def test_is_eq_precedence_flat(self): self.assert_compile( (self.table1.c.name == null()) @@ -4522,3 +4551,45 @@ class AnyAllTest(fixtures.TestBase, testing.AssertsCompiledSQL): r"use the .scalar_values\(\) method.", ): fn(values(t.c.data).data([(1,), (42,)])) + + +class BitOpTest(fixtures.TestBase, testing.AssertsCompiledSQL): + __dialect__ = "default" + + def test_compile_not_column_lvl(self): + c = column("c", Integer) + + self.assert_compile( + select(c.bitwise_not()), + "SELECT ~c", + ) + + def test_compile_not_colexpr_lvl(self): + c = column("c", Integer) + + self.assert_compile( + select(operators.bitwise_not_op(c)), + "SELECT ~c", + ) + + @testing.combinations( + ("xor", operators.bitwise_xor_op, "^"), + ("xor_lambda", lambda c1, c2: c1.bitwise_xor(c2), "^"), + ("or", operators.bitwise_or_op, "|"), + ("or_lambda", lambda c1, c2: c1.bitwise_or(c2), "|"), + ("and", operators.bitwise_and_op, "&"), + ("and_lambda", lambda c1, c2: c1.bitwise_and(c2), "&"), + ("lshift", operators.bitwise_lshift_op, "<<"), + ("ls_lambda", lambda c1, c2: c1.bitwise_lshift(c2), "<<"), + ("rshift", operators.bitwise_rshift_op, ">>"), + ("rs_lambda", lambda c1, c2: c1.bitwise_rshift(c2), ">>"), + id_="iaa", + ) + def test_compile_binary(self, py_op, sql_op): + c1 = column("c1", Integer) + c2 = column("c2", Integer) + + self.assert_compile( + select(py_op(c1, c2)), + f"SELECT c1 {sql_op} c2 AS anon_1", + ) -- 2.47.3