add
This commit is contained in:
macbook_max
+327
@@ -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__
|
||||
Reference in New Issue
Block a user