328 lines
10 KiB
Python
328 lines
10 KiB
Python
# 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__
|