Add binary operators for the BitStream class

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ default_install_hook_types:
   - post-checkout
   - post-merge
   - post-rewrite
-exclude: "(reports|logs|.devcontainer)/.*"
+exclude: "(firmware|reports|logs|.devcontainer|pymibs)/.*"
 repos:
   - repo: https://github.com/adrienverge/yamllint.git
     rev: v1.38.0
@@ -35,16 +35,14 @@ repos:
       - id: check-docstring-first
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.15.12
+    rev: v0.15.20
     hooks:
       # Run the linter.
       - id: ruff-check
-        name: ruff-check
         args: [--fix]
 
       # Run the formatter.
       - id: ruff-format
-        name: ruff-format
 
   - repo: local
     hooks:
@@ -56,12 +54,29 @@ repos:
         language: system
         entry: uv run mypy
         files: .*(\.py|\.pyi)$
-        args: [features, ci, scripts]
+        args: [src]
         pass_filenames: false
 
+  # - repo: local
+  #   hooks:
+  #     - id: pytest
+  #       name: pytest
+  #       stages:
+  #         - pre-commit
+  #         - pre-push
+  #       language: system
+  #       entry: uv run pytest
+  #       files: (src|tests)/.*(\.py|\.pyi)$
+  #       args:
+  #         - tests
+  #         - -q
+  #         - --no-summary
+  #         - --no-header
+  #       pass_filenames: false
+
   - repo: https://github.com/astral-sh/uv-pre-commit
     # uv version.
-    rev: 0.11.8
+    rev: 0.11.28
     hooks:
       # Ensure the lockfile is up-to-date
       - id: uv-lock

pyproject.toml

@@ -120,10 +120,8 @@ select = [
 ]
 fixable = ["ALL"]
 ignore = [
+    "D203",  # Incorrect black line before class
     "D100",  # Docstring in public module
-    "D102",  # Docstring in public method
-    "D203",  # Blank line before class
-    "D104",  # Docstring in public package
     "D105",  # Docstring in magic method
     "D212",  # Multiline Docstring Start
     "D205",  # Line between summary and description
@@ -133,9 +131,11 @@ ignore = [
 isort.split-on-trailing-comma = false
 
 [tool.ruff.lint.per-file-ignores]
+"**/__init__.py" = ["D104"]
 "tests/*.py" = ["SLF001", "PT013"]
-"typing.py" = ["D101", "D107"]
+"typing.py" = ["D101", "D107", "D102", "D100"]
 "exceptions.py" = ["D101", "D107"]
+"test_*.py" = ["D100", "D101", "D102"]
 
 [tool.ruff.format]
 # skip-magic-trailing-comma = true

src/bitstructures/base/bitstream.py

@@ -1,10 +1,11 @@
 from collections.abc import Generator
-from typing import Any, Self, overload
+from typing import Self, overload
 
 from bitarray import bitarray
 from bitarray.util import ba2int, int2ba
 
 from bitstructures.exceptions import ReadError, SizeError, WriteError
+from bitstructures.typing import Buffer, BufferCmp
 
 
 class BitStream:
@@ -18,7 +19,7 @@ class BitStream:
     instead of reading the memory object.
     """
 
-    def __init__(self, buffer: bitarray | bytes | str = b"", /) -> None:
+    def __init__(self, buffer: Buffer = b"", /) -> None:
         self.stream: bitarray
         assert isinstance(buffer, bitarray | bytes | str)
         if isinstance(buffer, bitarray):
@@ -82,13 +83,47 @@ def __bytes__(self) -> bytes:
         """
         if len(self) % 8 != 0:
             raise SizeError(
-                f"Cannot convert a BitStream of length {len(self)} to bytes, must be divisable by 8"
+                f"Cannot convert a {self.__class__.__name__} of length {len(self)} to bytes, "
+                f"must be divisable by 8"
             )
         return self.stream.tobytes()
 
     def __hash__(self) -> int:
         """Provide hashing functionality to our bitstream."""
-        return hash(self.stream)
+        return hash(self.stream.to01())
+
+    def __and__(self, other: BufferCmp) -> Self:
+        if isinstance(other, BitStream):
+            return self.__class__(self.stream & other.stream)
+        if isinstance(other, bitarray | bytes | str):
+            # Recurse into the above statement
+            return self & self.__class__(other)
+        return NotImplemented
+
+    def __or__(self, other: BufferCmp) -> Self:
+        if isinstance(other, BitStream):
+            return self.__class__(self.stream | other.stream)
+        if isinstance(other, bitarray | bytes | str):
+            # Recurse into the above statement
+            return self | self.__class__(other)
+        return NotImplemented
+
+    def __xor__(self, other: BufferCmp) -> Self:
+        if isinstance(other, BitStream):
+            return self.__class__(self.stream ^ other.stream)
+        if isinstance(other, bitarray | bytes | str):
+            # Recurse into the above statement
+            return self ^ self.__class__(other)
+        return NotImplemented
+
+    def __invert__(self) -> Self:
+        return self.__class__(~self.stream)
+
+    def __lshift__(self, n: int) -> Self:
+        return self.__class__(self.stream << n)
+
+    def __rshift__(self, n: int) -> Self:
+        return self.__class__(self.stream >> n)
 
     def __eq__(self, other: object) -> bool:
         """
@@ -97,71 +132,73 @@ def __eq__(self, other: object) -> bool:
         """
         if isinstance(other, BitStream):
             return self.stream == other.stream
-        if isinstance(other, bytes | str):
+        if isinstance(other, bitarray | bytes | str):
             # Recurse into the above statement
-            return self == BitStream(other)
-        raise TypeError(f"Cannot compare a {self.__class__.__name__} with a {type(other)}")
+            return self == self.__class__(other)
+        return False
 
-    def __add__(self, other: Any) -> "BitStream":
+    def __add__(self, other: BufferCmp) -> Self:
         """
         Append another string buffer into this buffer by writing the
         contents of the other buffer into our StringIO instance.
         """
         if isinstance(other, BitStream):
-            return BitStream(self.stream + other.stream)
+            return self.__class__(self.stream + other.stream)
         if isinstance(other, bitarray | bytes | str):
             # Recurse into the above statement
-            return self + BitStream(other)
-        raise TypeError(f"Cannot add {type(other)} to a {BitStream.__name__}")
+            return self + self.__class__(other)
+        return NotImplemented
 
-    def __iadd__(self, other: Any) -> Self:
+    def __iadd__(self, other: BufferCmp) -> Self:
         """
         Append another string buffer into this buffer by writing the
         contents of the other buffer into our StringIO instance.
         """
         if isinstance(other, BitStream):
             self.stream += other.stream
-        elif isinstance(other, bitarray | bytes | str):
-            self.stream += BitStream(other).stream
-        else:
-            raise TypeError(f"Cannot add {type(other)} to a {BitStream.__name__}")
-        return self
+            return self
+        if isinstance(other, bitarray | bytes | str):
+            # Recurse into the above statement
+            self.stream += self.__class__(other).stream
+            return self
+        return NotImplemented
 
     def __iter__(self) -> Generator[int]:
-        """Returns either a 1 or 0 for each bit in the stream."""
+        """Returns either a 1 or 0 for all bits in the stream."""
         yield from self.stream
 
     @overload
     def __getitem__(self, s: int) -> int: ...
     @overload
-    def __getitem__(self, s: slice) -> "BitStream": ...
-    def __getitem__(self, s: slice | int) -> "BitStream | int":
+    def __getitem__(self, s: slice) -> Self: ...
+    def __getitem__(self, s: slice | int) -> "Self | int":
         """Allow string slicing methods upon this class."""
         if isinstance(s, int):
+            # The bitarray library returns an integer when getitem is int
             return self.stream[s]
-        return BitStream(self.stream[s])
+        return self.__class__(self.stream[s])
 
-    def peek(self, size: int) -> "BitStream":
+    def peek(self, size: int) -> Self:
         """Looks through the string buffer without modifying the stream."""
-        return BitStream(self.stream[:size])
+        return self.__class__(self.stream[:size])
 
-    def read(self, size: int | None = -1) -> "BitStream":
+    def read(self, size: int | None = -1) -> Self:
         """Reads a specified amount of bits through the string buffer modifying the stream."""
         if size is None or size < 0:
             return self
         if size > len(self):
             raise ReadError(
                 f"Cannot read {size} bits, reached the EOS or "
-                f"the current stream is shorter that specified size. stream={len(self)}"
+                f"the current stream is shorter than specified size. stream={len(self)}"
             )
         try:
             # Modify the stream buffer
             out, self.stream = self.stream[:size], self.stream[size:]
-            return BitStream(out)
+            return self.__class__(out)
         except Exception as err:
             raise ReadError from err
 
-    def write(self, value: "int | BitStream", size: int) -> None:
+    def write(self, value: "int | Self", size: int) -> None:
         """
         Writes values or bit buffers of a specified bit size
         through the string buffer modifying the stream.
@@ -173,7 +210,13 @@ def write(self, value: "int | BitStream", size: int) -> None:
             )
         try:
             if isinstance(value, BitStream):
-                self += value
+                if value.bit_length() > size:
+                    raise WriteError(
+                        f"Cannot write a stream {value} of size {value.bit_length()}, "
+                        f"into a stream of size {size}"
+                    )
+                # Left pad the stream with bits
+                self += value.ljust(size)
                 return
             self += int2ba(value, size)
         except Exception as err:
@@ -186,3 +229,21 @@ def copy(self) -> Self:
     def bit_length(self) -> int:
         """Just returns the len of this BitStream, meant to mirror int.bit_length()."""
         return len(self)
+
+    def ljust(self, width: int, fillbit: int = 0) -> Self:
+        """Returns a copy of the stream right-padded with the fill character."""
+        if fillbit not in {0, 1}:
+            raise ValueError("fillbit must be '0' or '1'")
+        if len(self) >= width:
+            return self.copy()
+        padding = bitarray([fillbit]) * (width - len(self))
+        return self + padding
+
+    def rjust(self, width: int, fillbit: int = 0) -> Self:
+        """Returns a copy of the stream left-padded with the fill character."""
+        if fillbit not in {0, 1}:
+            raise ValueError("fillbit must be '0' or '1'")
+        if len(self) >= width:
+            return self.copy()
+        padding = bitarray([fillbit]) * (width - len(self))
+        return self.__class__(padding + self.stream)

src/bitstructures/base/codec.py

@@ -47,6 +47,7 @@ class SingletonMeta(__protocol_type):  # type: ignore[misc, valid-type]
     _instances: ClassVar[dict["SingletonMeta", type]] = {}
 
     def __call__(self, *args: Any, **kwargs: Any) -> type:
+        """Returns an instance of this class upon calling again."""
         if self not in self._instances:
             self._instances[self] = super().__call__(*args, **kwargs)
         return self._instances[self]
@@ -69,6 +70,7 @@ def _validate_item(self, item: "Codec | StackC") -> None:
         return
 
     def pprint(self, *, depth: int = 1) -> str:
+        """Returns a nice representation of this class with padding."""
         # RECURSIVE
         padding = "\t{t}{v:-^38}{t}\n"
         stack = ""
@@ -128,6 +130,7 @@ def __init__(self, subcodec: "Codec | None" = None) -> None:
 
     @property
     def name(self) -> str:
+        """Returns the name of this codec, set via 'name' / Codec."""
         return self._name
 
     @property
@@ -150,6 +153,7 @@ def size(self) -> int:
 
     @property
     def subcodec(self) -> "Codec":
+        """Returns the Subcodec that this Codec is based on."""
         return self._subcodec or self
 
     def _check_initialized(self) -> None:
@@ -398,6 +402,7 @@ class Struct(Codec, StructProtocol):
 
     @property
     def subcodecs(self) -> StackC:
+        """Returns the stack of Codec's listed in this Struct."""
         return self._subcodec_stack
 
     @override
@@ -574,6 +579,7 @@ def build(self, container: Container) -> BitStream:
             raise BuildError(io, container, codecs, repr(err)) from err
 
     def sizeof(self, io: BitStream, context: Container, codecs: StackC) -> int:
+        """Returns the bit-size of this Structure/Codec."""
         if isinstance(self.size, int) and self.size > 0:
             return self.size
         return super().sizeof(io, context, codecs)
@@ -1020,6 +1026,7 @@ def __init__(
 
     @property
     def enum(self) -> EnumBase:
+        """Returns a read only enumeration object created via this class."""
         return self._enum
 
     @override