SQLAlchemy 2.1.0b2__cp313-cp313t-win_arm64.whl

sqlalchemy/dialects/postgresql/bitstring.py

@@ -0,0 +1,327 @@
+# dialects/postgresql/bitstring.py
+# Copyright (C) 2013-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+from __future__ import annotations
+
+import math
+from typing import Any
+from typing import cast
+from typing import Literal
+from typing import SupportsIndex
+
+
+class BitString(str):
+    """Represent a PostgreSQL bit string in python.
+
+    This object is used by the :class:`_postgresql.BIT` type when returning
+    values.   :class:`_postgresql.BitString` values may also be constructed
+    directly and used with :class:`_postgresql.BIT` columns::
+
+        from sqlalchemy.dialects.postgresql import BitString
+
+        with engine.connect() as conn:
+            conn.execute(table.insert(), {"data": BitString("011001101")})
+
+    .. versionadded:: 2.1
+
+    """
+
+    _DIGITS = frozenset("01")
+
+    def __new__(cls, _value: str, _check: bool = True) -> BitString:
+        if isinstance(_value, BitString):
+            return _value
+        elif _check and cls._DIGITS.union(_value) > cls._DIGITS:
+            raise ValueError("BitString must only contain '0' and '1' chars")
+        else:
+            return super().__new__(cls, _value)
+
+    @classmethod
+    def from_int(cls, value: int, length: int) -> BitString:
+        """Returns a BitString consisting of the bits in the integer ``value``.
+        A ``ValueError`` is raised if ``value`` is not a non-negative integer.
+
+        If the provided ``value`` can not be represented in a bit string
+        of at most ``length``, a ``ValueError`` will be raised. The bitstring
+        will be padded on the left by ``'0'`` to bits to produce a
+        bitstring of the desired length.
+        """
+        if value < 0:
+            raise ValueError("value must be non-negative")
+        if length < 0:
+            raise ValueError("length must be non-negative")
+
+        template_str = f"{{0:0{length}b}}" if length > 0 else ""
+        r = template_str.format(value)
+
+        if (length == 0 and value > 0) or len(r) > length:
+            raise ValueError(
+                f"Cannot encode {value} as a BitString of length {length}"
+            )
+
+        return cls(r)
+
+    @classmethod
+    def from_bytes(cls, value: bytes, length: int = -1) -> BitString:
+        """Returns a ``BitString`` consisting of the bits in the given
+        ``value`` bytes.
+
+        If ``length`` is provided, then the length of the provided string
+        will be exactly ``length``, with ``'0'`` bits inserted at the left of
+        the string in order to produce a value of the required length.
+        If the bits obtained by omitting the leading ``'0'`` bits of ``value``
+        cannot be represented in a string of this length a ``ValueError``
+        will be raised.
+        """
+        str_v: str = "".join(f"{int(c):08b}" for c in value)
+        if length >= 0:
+            str_v = str_v.lstrip("0")
+
+            if len(str_v) > length:
+                raise ValueError(
+                    f"Cannot encode {value!r} as a BitString of "
+                    f"length {length}"
+                )
+            str_v = str_v.zfill(length)
+
+        return cls(str_v)
+
+    def get_bit(self, index: int) -> Literal["0", "1"]:
+        """Returns the value of the flag at the given
+        index::
+
+            BitString("0101").get_flag(4) == "1"
+        """
+        return cast(Literal["0", "1"], super().__getitem__(index))
+
+    @property
+    def bit_length(self) -> int:
+        return len(self)
+
+    @property
+    def octet_length(self) -> int:
+        return math.ceil(len(self) / 8)
+
+    def has_bit(self, index: int) -> bool:
+        return self.get_bit(index) == "1"
+
+    def set_bit(
+        self, index: int, value: bool | int | Literal["0", "1"]
+    ) -> BitString:
+        """Set the bit at index to the given value.
+
+        If value is an int, then it is considered to be '1' iff nonzero.
+        """
+        if index < 0 or index >= len(self):
+            raise IndexError("BitString index out of range")
+
+        if isinstance(value, (bool, int)):
+            value = "1" if value else "0"
+
+        if self.get_bit(index) == value:
+            return self
+
+        return BitString(
+            "".join([self[:index], value, self[index + 1 :]]), False
+        )
+
+    def lstrip(self, char: str | None = None) -> BitString:
+        """Returns a copy of the BitString with leading characters removed.
+
+        If omitted or None, 'chars' defaults '0'::
+
+            BitString("00010101000").lstrip() == BitString("00010101")
+            BitString("11110101111").lstrip("1") == BitString("1111010")
+        """
+        if char is None:
+            char = "0"
+        return BitString(super().lstrip(char), False)
+
+    def rstrip(self, char: str | None = "0") -> BitString:
+        """Returns a copy of the BitString with trailing characters removed.
+
+        If omitted or None, ``'char'`` defaults to "0"::
+
+            BitString("00010101000").rstrip() == BitString("10101000")
+            BitString("11110101111").rstrip("1") == BitString("10101111")
+        """
+        if char is None:
+            char = "0"
+        return BitString(super().rstrip(char), False)
+
+    def strip(self, char: str | None = "0") -> BitString:
+        """Returns a copy of the BitString with both leading and trailing
+        characters removed.
+        If omitted or None, ``'char'`` defaults to ``"0"``::
+
+            BitString("00010101000").rstrip() == BitString("10101")
+            BitString("11110101111").rstrip("1") == BitString("1010")
+        """
+        if char is None:
+            char = "0"
+        return BitString(super().strip(char))
+
+    def removeprefix(self, prefix: str, /) -> BitString:
+        return BitString(super().removeprefix(prefix), False)
+
+    def removesuffix(self, suffix: str, /) -> BitString:
+        return BitString(super().removesuffix(suffix), False)
+
+    def replace(
+        self,
+        old: str,
+        new: str,
+        count: SupportsIndex = -1,
+    ) -> BitString:
+        new = BitString(new)
+        return BitString(super().replace(old, new, count), False)
+
+    def split(
+        self,
+        sep: str | None = None,
+        maxsplit: SupportsIndex = -1,
+    ) -> list[str]:
+        return [BitString(word) for word in super().split(sep, maxsplit)]
+
+    def zfill(self, width: SupportsIndex) -> BitString:
+        return BitString(super().zfill(width), False)
+
+    def __repr__(self) -> str:
+        return f'BitString("{self.__str__()}")'
+
+    def __int__(self) -> int:
+        return int(self, 2) if self else 0
+
+    def to_bytes(self, length: int = -1) -> bytes:
+        return int(self).to_bytes(
+            length if length >= 0 else self.octet_length, byteorder="big"
+        )
+
+    def __bytes__(self) -> bytes:
+        return self.to_bytes()
+
+    def __getitem__(
+        self, key: SupportsIndex | slice[Any, Any, Any]
+    ) -> BitString:
+        return BitString(super().__getitem__(key), False)
+
+    def __add__(self, o: str) -> BitString:
+        """Return self + o"""
+        if not isinstance(o, str):
+            raise TypeError(
+                f"Can only concatenate str (not '{type(self)}') to BitString"
+            )
+        return BitString("".join([self, o]))
+
+    def __radd__(self, o: str) -> BitString:
+        if not isinstance(o, str):
+            raise TypeError(
+                f"Can only concatenate str (not '{type(self)}') to BitString"
+            )
+        return BitString("".join([o, self]))
+
+    def __lshift__(self, amount: int) -> BitString:
+        """Shifts each the bitstring to the left by the given amount.
+        String length is preserved::
+
+            BitString("000101") << 1 == BitString("001010")
+        """
+        return BitString(
+            "".join([self, *("0" for _ in range(amount))])[-len(self) :], False
+        )
+
+    def __rshift__(self, amount: int) -> BitString:
+        """Shifts each bit in the bitstring to the right by the given amount.
+        String length is preserved::
+
+            BitString("101") >> 1 == BitString("010")
+        """
+        return BitString(self[:-amount], False).zfill(width=len(self))
+
+    def __invert__(self) -> BitString:
+        """Inverts (~) each bit in the
+        bitstring::
+
+            ~BitString("01010") == BitString("10101")
+        """
+        return BitString("".join("1" if x == "0" else "0" for x in self))
+
+    def __and__(self, o: str) -> BitString:
+        """Performs a bitwise and (``&``) with the given operand.
+        A ``ValueError`` is raised if the operand is not the same length.
+
+        e.g.::
+
+            BitString("011") & BitString("011") == BitString("010")
+        """
+
+        if not isinstance(o, str):
+            return NotImplemented
+        o = BitString(o)
+        if len(self) != len(o):
+            raise ValueError("Operands must be the same length")
+
+        return BitString(
+            "".join(
+                "1" if (x == "1" and y == "1") else "0"
+                for x, y in zip(self, o)
+            ),
+            False,
+        )
+
+    def __or__(self, o: str) -> BitString:
+        """Performs a bitwise or (``|``) with the given operand.
+        A ``ValueError`` is raised if the operand is not the same length.
+
+        e.g.::
+
+            BitString("011") | BitString("010") == BitString("011")
+        """
+        if not isinstance(o, str):
+            return NotImplemented
+
+        if len(self) != len(o):
+            raise ValueError("Operands must be the same length")
+
+        o = BitString(o)
+        return BitString(
+            "".join(
+                "1" if (x == "1" or y == "1") else "0"
+                for (x, y) in zip(self, o)
+            ),
+            False,
+        )
+
+    def __xor__(self, o: str) -> BitString:
+        """Performs a bitwise xor (``^``) with the given operand.
+        A ``ValueError`` is raised if the operand is not the same length.
+
+        e.g.::
+
+            BitString("011") ^ BitString("010") == BitString("001")
+        """
+
+        if not isinstance(o, BitString):
+            return NotImplemented
+
+        if len(self) != len(o):
+            raise ValueError("Operands must be the same length")
+
+        return BitString(
+            "".join(
+                (
+                    "1"
+                    if ((x == "1" and y == "0") or (x == "0" and y == "1"))
+                    else "0"
+                )
+                for (x, y) in zip(self, o)
+            ),
+            False,
+        )
+
+    __rand__ = __and__
+    __ror__ = __or__
+    __rxor__ = __xor__

sqlalchemy/dialects/postgresql/dml.py

@@ -0,0 +1,360 @@
+# dialects/postgresql/dml.py
+# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+from __future__ import annotations
+
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+from . import ext
+from .._typing import _OnConflictConstraintT
+from .._typing import _OnConflictIndexElementsT
+from .._typing import _OnConflictIndexWhereT
+from .._typing import _OnConflictSetT
+from .._typing import _OnConflictWhereT
+from ... import util
+from ...sql import coercions
+from ...sql import roles
+from ...sql import schema
+from ...sql._typing import _DMLTableArgument
+from ...sql.base import _exclusive_against
+from ...sql.base import ColumnCollection
+from ...sql.base import ReadOnlyColumnCollection
+from ...sql.base import SyntaxExtension
+from ...sql.dml import _DMLColumnElement
+from ...sql.dml import Insert as StandardInsert
+from ...sql.elements import ClauseElement
+from ...sql.elements import ColumnElement
+from ...sql.elements import KeyedColumnElement
+from ...sql.elements import TextClause
+from ...sql.expression import alias
+from ...sql.type_api import NULLTYPE
+from ...sql.visitors import InternalTraversal
+from ...util.typing import Self
+
+__all__ = ("Insert", "insert")
+
+
+def insert(table: _DMLTableArgument) -> Insert:
+    """Construct a PostgreSQL-specific variant :class:`_postgresql.Insert`
+    construct.
+
+    .. container:: inherited_member
+
+        The :func:`sqlalchemy.dialects.postgresql.insert` function creates
+        a :class:`sqlalchemy.dialects.postgresql.Insert`.  This class is based
+        on the dialect-agnostic :class:`_sql.Insert` construct which may
+        be constructed using the :func:`_sql.insert` function in
+        SQLAlchemy Core.
+
+    The :class:`_postgresql.Insert` construct includes additional methods
+    :meth:`_postgresql.Insert.on_conflict_do_update`,
+    :meth:`_postgresql.Insert.on_conflict_do_nothing`.
+
+    """
+    return Insert(table)
+
+
+class Insert(StandardInsert):
+    """PostgreSQL-specific implementation of INSERT.
+
+    Adds methods for PG-specific syntaxes such as ON CONFLICT.
+
+    The :class:`_postgresql.Insert` object is created using the
+    :func:`sqlalchemy.dialects.postgresql.insert` function.
+
+    """
+
+    stringify_dialect = "postgresql"
+    inherit_cache = True
+
+    @util.memoized_property
+    def excluded(
+        self,
+    ) -> ReadOnlyColumnCollection[str, KeyedColumnElement[Any]]:
+        """Provide the ``excluded`` namespace for an ON CONFLICT statement
+
+        PG's ON CONFLICT clause allows reference to the row that would
+        be inserted, known as ``excluded``.  This attribute provides
+        all columns in this row to be referenceable.
+
+        .. tip::  The :attr:`_postgresql.Insert.excluded` attribute is an
+            instance of :class:`_expression.ColumnCollection`, which provides
+            an interface the same as that of the :attr:`_schema.Table.c`
+            collection described at :ref:`metadata_tables_and_columns`.
+            With this collection, ordinary names are accessible like attributes
+            (e.g. ``stmt.excluded.some_column``), but special names and
+            dictionary method names should be accessed using indexed access,
+            such as ``stmt.excluded["column name"]`` or
+            ``stmt.excluded["values"]``.   See the docstring for
+            :class:`_expression.ColumnCollection` for further examples.
+
+        .. seealso::
+
+            :ref:`postgresql_insert_on_conflict` - example of how
+            to use :attr:`_expression.Insert.excluded`
+
+        """
+        return alias(self.table, name="excluded").columns
+
+    _on_conflict_exclusive = _exclusive_against(
+        "_post_values_clause",
+        msgs={
+            "_post_values_clause": "This Insert construct already has "
+            "an ON CONFLICT clause established"
+        },
+    )
+
+    @_on_conflict_exclusive
+    def on_conflict_do_update(
+        self,
+        constraint: _OnConflictConstraintT = None,
+        index_elements: _OnConflictIndexElementsT = None,
+        index_where: _OnConflictIndexWhereT = None,
+        set_: _OnConflictSetT = None,
+        where: _OnConflictWhereT = None,
+    ) -> Self:
+        r"""
+        Specifies a DO UPDATE SET action for ON CONFLICT clause.
+
+        Either the ``constraint`` or ``index_elements`` argument is
+        required, but only one of these can be specified.
+
+        :param constraint:
+         The name of a unique or exclusion constraint on the table,
+         or the constraint object itself if it has a .name attribute.
+
+        :param index_elements:
+         A sequence consisting of string column names, :class:`_schema.Column`
+         objects, or other column expression objects that will be used
+         to infer a target index.
+
+        :param index_where:
+         Additional WHERE criterion that can be used to infer a
+         conditional target index.
+
+        :param set\_:
+         A dictionary or other mapping object
+         where the keys are either names of columns in the target table,
+         or :class:`_schema.Column` objects or other ORM-mapped columns
+         matching that of the target table, and expressions or literals
+         as values, specifying the ``SET`` actions to take.
+
+         .. versionadded:: 1.4 The
+            :paramref:`_postgresql.Insert.on_conflict_do_update.set_`
+            parameter supports :class:`_schema.Column` objects from the target
+            :class:`_schema.Table` as keys.
+
+         .. warning:: This dictionary does **not** take into account
+            Python-specified default UPDATE values or generation functions,
+            e.g. those specified using :paramref:`_schema.Column.onupdate`.
+            These values will not be exercised for an ON CONFLICT style of
+            UPDATE, unless they are manually specified in the
+            :paramref:`.Insert.on_conflict_do_update.set_` dictionary.
+
+        :param where:
+         Optional argument. An expression object representing a ``WHERE``
+         clause that restricts the rows affected by ``DO UPDATE SET``. Rows not
+         meeting the ``WHERE`` condition will not be updated (effectively a
+         ``DO NOTHING`` for those rows).
+
+
+        .. seealso::
+
+            :ref:`postgresql_insert_on_conflict`
+
+        """
+        return self.ext(
+            OnConflictDoUpdate(
+                constraint, index_elements, index_where, set_, where
+            )
+        )
+
+    @_on_conflict_exclusive
+    def on_conflict_do_nothing(
+        self,
+        constraint: _OnConflictConstraintT = None,
+        index_elements: _OnConflictIndexElementsT = None,
+        index_where: _OnConflictIndexWhereT = None,
+    ) -> Self:
+        """
+        Specifies a DO NOTHING action for ON CONFLICT clause.
+
+        The ``constraint`` and ``index_elements`` arguments
+        are optional, but only one of these can be specified.
+
+        :param constraint:
+         The name of a unique or exclusion constraint on the table,
+         or the constraint object itself if it has a .name attribute.
+
+        :param index_elements:
+         A sequence consisting of string column names, :class:`_schema.Column`
+         objects, or other column expression objects that will be used
+         to infer a target index.
+
+        :param index_where:
+         Additional WHERE criterion that can be used to infer a
+         conditional target index.
+
+        .. seealso::
+
+            :ref:`postgresql_insert_on_conflict`
+
+        """
+        return self.ext(
+            OnConflictDoNothing(constraint, index_elements, index_where)
+        )
+
+
+class OnConflictClause(SyntaxExtension, ClauseElement):
+    stringify_dialect = "postgresql"
+
+    constraint_target: Optional[str]
+    inferred_target_elements: Optional[List[Union[str, schema.Column[Any]]]]
+    inferred_target_whereclause: Optional[
+        Union[ColumnElement[Any], TextClause]
+    ]
+
+    _traverse_internals = [
+        ("constraint_target", InternalTraversal.dp_string),
+        ("inferred_target_elements", InternalTraversal.dp_multi_list),
+        ("inferred_target_whereclause", InternalTraversal.dp_clauseelement),
+    ]
+
+    def __init__(
+        self,
+        constraint: _OnConflictConstraintT = None,
+        index_elements: _OnConflictIndexElementsT = None,
+        index_where: _OnConflictIndexWhereT = None,
+    ):
+        if constraint is not None:
+            if not isinstance(constraint, str) and isinstance(
+                constraint,
+                (schema.Constraint, ext.ExcludeConstraint),
+            ):
+                constraint = getattr(constraint, "name") or constraint
+
+        if constraint is not None:
+            if index_elements is not None:
+                raise ValueError(
+                    "'constraint' and 'index_elements' are mutually exclusive"
+                )
+
+            if isinstance(constraint, str):
+                self.constraint_target = constraint
+                self.inferred_target_elements = None
+                self.inferred_target_whereclause = None
+            elif isinstance(constraint, schema.Index):
+                index_elements = constraint.expressions
+                index_where = constraint.dialect_options["postgresql"].get(
+                    "where"
+                )
+            elif isinstance(constraint, ext.ExcludeConstraint):
+                index_elements = constraint.columns
+                index_where = constraint.where
+            else:
+                index_elements = constraint.columns
+                index_where = constraint.dialect_options["postgresql"].get(
+                    "where"
+                )
+
+        if index_elements is not None:
+            self.constraint_target = None
+            self.inferred_target_elements = [
+                coercions.expect(roles.DDLConstraintColumnRole, column)
+                for column in index_elements
+            ]
+
+            self.inferred_target_whereclause = (
+                coercions.expect(
+                    (
+                        roles.StatementOptionRole
+                        if isinstance(constraint, ext.ExcludeConstraint)
+                        else roles.WhereHavingRole
+                    ),
+                    index_where,
+                )
+                if index_where is not None
+                else None
+            )
+
+        elif constraint is None:
+            self.constraint_target = self.inferred_target_elements = (
+                self.inferred_target_whereclause
+            ) = None
+
+    def apply_to_insert(self, insert_stmt: StandardInsert) -> None:
+        insert_stmt.apply_syntax_extension_point(
+            self.append_replacing_same_type, "post_values"
+        )
+
+
+class OnConflictDoNothing(OnConflictClause):
+    __visit_name__ = "on_conflict_do_nothing"
+
+    inherit_cache = True
+
+
+class OnConflictDoUpdate(OnConflictClause):
+    __visit_name__ = "on_conflict_do_update"
+
+    update_values_to_set: Dict[_DMLColumnElement, ColumnElement[Any]]
+    update_whereclause: Optional[ColumnElement[Any]]
+
+    _traverse_internals = OnConflictClause._traverse_internals + [
+        ("update_values_to_set", InternalTraversal.dp_dml_values),
+        ("update_whereclause", InternalTraversal.dp_clauseelement),
+    ]
+
+    def __init__(
+        self,
+        constraint: _OnConflictConstraintT = None,
+        index_elements: _OnConflictIndexElementsT = None,
+        index_where: _OnConflictIndexWhereT = None,
+        set_: _OnConflictSetT = None,
+        where: _OnConflictWhereT = None,
+    ):
+        super().__init__(
+            constraint=constraint,
+            index_elements=index_elements,
+            index_where=index_where,
+        )
+
+        if (
+            self.inferred_target_elements is None
+            and self.constraint_target is None
+        ):
+            raise ValueError(
+                "Either constraint or index_elements, "
+                "but not both, must be specified unless DO NOTHING"
+            )
+
+        if isinstance(set_, dict):
+            if not set_:
+                raise ValueError("set parameter dictionary must not be empty")
+        elif isinstance(set_, ColumnCollection):
+            set_ = dict(set_)
+        else:
+            raise ValueError(
+                "set parameter must be a non-empty dictionary "
+                "or a ColumnCollection such as the `.c.` collection "
+                "of a Table object"
+            )
+
+        self.update_values_to_set = {
+            coercions.expect(roles.DMLColumnRole, k): coercions.expect(
+                roles.ExpressionElementRole, v, type_=NULLTYPE, is_crud=True
+            )
+            for k, v in set_.items()
+        }
+        self.update_whereclause = (
+            coercions.expect(roles.WhereHavingRole, where)
+            if where is not None
+            else None
+        )