1from __future__ import annotations
2
3import copy
4import numbers
5import re
6from collections import abc
7from typing import Union, List, Iterable, Any, Optional
8from bitstring import utils
9from bitstring.exceptions import CreationError, Error
10from bitstring.bits import Bits, BitsType, TBits
11
12import bitstring.dtypes
13
14class BitArray(Bits):
15 """A container holding a mutable sequence of bits.
16
17 Subclass of the immutable Bits class. Inherits all of its
18 methods (except __hash__) and adds mutating methods.
19
20 Mutating methods:
21
22 append() -- Append a bitstring.
23 byteswap() -- Change byte endianness in-place.
24 clear() -- Remove all bits from the bitstring.
25 insert() -- Insert a bitstring.
26 invert() -- Flip bit(s) between one and zero.
27 overwrite() -- Overwrite a section with a new bitstring.
28 prepend() -- Prepend a bitstring.
29 replace() -- Replace occurrences of one bitstring with another.
30 reverse() -- Reverse bits in-place.
31 rol() -- Rotate bits to the left.
32 ror() -- Rotate bits to the right.
33 set() -- Set bit(s) to 1 or 0.
34
35 Methods inherited from Bits:
36
37 all() -- Check if all specified bits are set to 1 or 0.
38 any() -- Check if any of specified bits are set to 1 or 0.
39 copy() -- Return a copy of the bitstring.
40 count() -- Count the number of bits set to 1 or 0.
41 cut() -- Create generator of constant sized chunks.
42 endswith() -- Return whether the bitstring ends with a sub-string.
43 find() -- Find a sub-bitstring in the current bitstring.
44 findall() -- Find all occurrences of a sub-bitstring in the current bitstring.
45 fromstring() -- Create a bitstring from a formatted string.
46 join() -- Join bitstrings together using current bitstring.
47 pp() -- Pretty print the bitstring.
48 rfind() -- Seek backwards to find a sub-bitstring.
49 split() -- Create generator of chunks split by a delimiter.
50 startswith() -- Return whether the bitstring starts with a sub-bitstring.
51 tobitarray() -- Return bitstring as a bitarray from the bitarray package.
52 tobytes() -- Return bitstring as bytes, padding if needed.
53 tofile() -- Write bitstring to file, padding if needed.
54 unpack() -- Interpret bits using format string.
55
56 Special methods:
57
58 Mutating operators are available: [], <<=, >>=, +=, *=, &=, |= and ^=
59 in addition to the inherited [], ==, !=, +, *, ~, <<, >>, &, | and ^.
60
61 Properties:
62
63 [GENERATED_PROPERTY_DESCRIPTIONS]
64
65 len -- Length of the bitstring in bits.
66
67 """
68
69 __slots__ = ()
70
71 # As BitArray objects are mutable, we shouldn't allow them to be hashed.
72 __hash__: None = None
73
74
75 def __init__(self, auto: Optional[Union[BitsType, int]] = None, /, length: Optional[int] = None,
76 offset: Optional[int] = None, **kwargs) -> None:
77 """Either specify an 'auto' initialiser:
78 A string of comma separated tokens, an integer, a file object,
79 a bytearray, a boolean iterable or another bitstring.
80
81 Or initialise via **kwargs with one (and only one) of:
82 bin -- binary string representation, e.g. '0b001010'.
83 hex -- hexadecimal string representation, e.g. '0x2ef'
84 oct -- octal string representation, e.g. '0o777'.
85 bytes -- raw data as a bytes object, for example read from a binary file.
86 int -- a signed integer.
87 uint -- an unsigned integer.
88 float / floatbe -- a big-endian floating point number.
89 bool -- a boolean (True or False).
90 se -- a signed exponential-Golomb code.
91 ue -- an unsigned exponential-Golomb code.
92 sie -- a signed interleaved exponential-Golomb code.
93 uie -- an unsigned interleaved exponential-Golomb code.
94 floatle -- a little-endian floating point number.
95 floatne -- a native-endian floating point number.
96 bfloat / bfloatbe - a big-endian bfloat format 16-bit floating point number.
97 bfloatle -- a little-endian bfloat format 16-bit floating point number.
98 bfloatne -- a native-endian bfloat format 16-bit floating point number.
99 intbe -- a signed big-endian whole byte integer.
100 intle -- a signed little-endian whole byte integer.
101 intne -- a signed native-endian whole byte integer.
102 uintbe -- an unsigned big-endian whole byte integer.
103 uintle -- an unsigned little-endian whole byte integer.
104 uintne -- an unsigned native-endian whole byte integer.
105 filename -- the path of a file which will be opened in binary read-only mode.
106
107 Other keyword arguments:
108 length -- length of the bitstring in bits, if needed and appropriate.
109 It must be supplied for all integer and float initialisers.
110 offset -- bit offset to the data. These offset bits are
111 ignored and this is intended for use when
112 initialising using 'bytes' or 'filename'.
113
114 """
115 if self._bitstore.immutable:
116 self._bitstore = self._bitstore._copy()
117 self._bitstore.immutable = False
118
119 def copy(self: TBits) -> TBits:
120 """Return a copy of the bitstring."""
121 return self.__copy__()
122
123 def __setattr__(self, attribute, value) -> None:
124 try:
125 # First try the ordinary attribute setter
126 super().__setattr__(attribute, value)
127 except AttributeError:
128 dtype = bitstring.dtypes.Dtype(attribute)
129 x = object.__new__(Bits)
130 if (set_fn := dtype.set_fn) is None:
131 raise AttributeError(f"Cannot set attribute '{attribute}' as it does not have a set_fn.")
132 set_fn(x, value)
133 if len(x) != dtype.bitlength:
134 raise CreationError(f"Can't initialise with value of length {len(x)} bits, "
135 f"as attribute has length of {dtype.bitlength} bits.")
136 self._bitstore = x._bitstore
137 return
138
139 def __iadd__(self, bs: BitsType) -> BitArray:
140 """Append bs to current bitstring. Return self.
141
142 bs -- the bitstring to append.
143
144 """
145 self._append(bs)
146 return self
147
148 def __copy__(self) -> BitArray:
149 """Return a new copy of the BitArray."""
150 s_copy = BitArray()
151 s_copy._bitstore = self._bitstore._copy()
152 assert s_copy._bitstore.immutable is False
153 return s_copy
154
155 def _setitem_int(self, key: int, value: Union[BitsType, int]) -> None:
156 if isinstance(value, numbers.Integral):
157 if value == 0:
158 self._bitstore[key] = 0
159 return
160 if value in (1, -1):
161 self._bitstore[key] = 1
162 return
163 raise ValueError(f"Cannot set a single bit with integer {value}.")
164 try:
165 value = self._create_from_bitstype(value)
166 except TypeError:
167 raise TypeError(f"Bitstring, integer or string expected. Got {type(value)}.")
168 positive_key = key + len(self) if key < 0 else key
169 if positive_key < 0 or positive_key >= len(self._bitstore):
170 raise IndexError(f"Bit position {key} out of range.")
171 self._bitstore[positive_key: positive_key + 1] = value._bitstore
172
173 def _setitem_slice(self, key: slice, value: BitsType) -> None:
174 if isinstance(value, numbers.Integral):
175 if key.step not in [None, -1, 1]:
176 if value in [0, 1]:
177 self.set(value, range(*key.indices(len(self))))
178 return
179 else:
180 raise ValueError("Can't assign an integer except 0 or 1 to a slice with a step value.")
181 # To find the length we first get the slice
182 s = self._bitstore.getslice(key.start, key.stop)
183 length = len(s)
184 # Now create an int of the correct length
185 if value >= 0:
186 value = self.__class__(uint=value, length=length)
187 else:
188 value = self.__class__(int=value, length=length)
189 else:
190 try:
191 value = self._create_from_bitstype(value)
192 except TypeError:
193 raise TypeError(f"Bitstring, integer or string expected. Got {type(value)}.")
194 self._bitstore.__setitem__(key, value._bitstore)
195
196 def __setitem__(self, key: Union[slice, int], value: BitsType) -> None:
197 if isinstance(key, numbers.Integral):
198 self._setitem_int(key, value)
199 else:
200 self._setitem_slice(key, value)
201
202 def __delitem__(self, key: Union[slice, int]) -> None:
203 """Delete item or range.
204
205 >>> a = BitArray('0x001122')
206 >>> del a[8:16]
207 >>> print a
208 0x0022
209
210 """
211 self._bitstore.__delitem__(key)
212 return
213
214 def __ilshift__(self: TBits, n: int) -> TBits:
215 """Shift bits by n to the left in place. Return self.
216
217 n -- the number of bits to shift. Must be >= 0.
218
219 """
220 if n < 0:
221 raise ValueError("Cannot shift by a negative amount.")
222 if not len(self):
223 raise ValueError("Cannot shift an empty bitstring.")
224 if not n:
225 return self
226 n = min(n, len(self))
227 return self._ilshift(n)
228
229 def __irshift__(self: TBits, n: int) -> TBits:
230 """Shift bits by n to the right in place. Return self.
231
232 n -- the number of bits to shift. Must be >= 0.
233
234 """
235 if n < 0:
236 raise ValueError("Cannot shift by a negative amount.")
237 if not len(self):
238 raise ValueError("Cannot shift an empty bitstring.")
239 if not n:
240 return self
241 n = min(n, len(self))
242 return self._irshift(n)
243
244 def __imul__(self: TBits, n: int) -> TBits:
245 """Concatenate n copies of self in place. Return self.
246
247 Called for expressions of the form 'a *= 3'.
248 n -- The number of concatenations. Must be >= 0.
249
250 """
251 if n < 0:
252 raise ValueError("Cannot multiply by a negative integer.")
253 return self._imul(n)
254
255 def __ior__(self: TBits, bs: BitsType) -> TBits:
256 bs = self._create_from_bitstype(bs)
257 self._bitstore |= bs._bitstore
258 return self
259
260 def __iand__(self: TBits, bs: BitsType) -> TBits:
261 bs = self._create_from_bitstype(bs)
262 self._bitstore &= bs._bitstore
263 return self
264
265 def __ixor__(self: TBits, bs: BitsType) -> TBits:
266 bs = self._create_from_bitstype(bs)
267 self._bitstore ^= bs._bitstore
268 return self
269
270 def _replace(self, old: Bits, new: Bits, start: int, end: int, count: int, bytealigned: Optional[bool]) -> int:
271 if bytealigned is None:
272 bytealigned = bitstring.options.bytealigned
273 # First find all the places where we want to do the replacements
274 starting_points: List[int] = []
275 for x in self.findall(old, start, end, bytealigned=bytealigned):
276 if not starting_points:
277 starting_points.append(x)
278 elif x >= starting_points[-1] + len(old):
279 # Can only replace here if it hasn't already been replaced!
280 starting_points.append(x)
281 if count != 0 and len(starting_points) == count:
282 break
283 if not starting_points:
284 return 0
285 replacement_list = [self._bitstore.getslice(0, starting_points[0])]
286 for i in range(len(starting_points) - 1):
287 replacement_list.append(new._bitstore)
288 replacement_list.append(
289 self._bitstore.getslice(starting_points[i] + len(old), starting_points[i + 1]))
290 # Final replacement
291 replacement_list.append(new._bitstore)
292 replacement_list.append(self._bitstore.getslice(starting_points[-1] + len(old), None))
293 if bitstring.options.lsb0:
294 # Addition of bitarray is always on the right, so assemble from other end
295 replacement_list.reverse()
296 self._bitstore.clear()
297 for r in replacement_list:
298 self._bitstore += r
299 return len(starting_points)
300
301 def replace(self, old: BitsType, new: BitsType, start: Optional[int] = None, end: Optional[int] = None,
302 count: Optional[int] = None, bytealigned: Optional[bool] = None) -> int:
303 """Replace all occurrences of old with new in place.
304
305 Returns number of replacements made.
306
307 old -- The bitstring to replace.
308 new -- The replacement bitstring.
309 start -- Any occurrences that start before this will not be replaced.
310 Defaults to 0.
311 end -- Any occurrences that finish after this will not be replaced.
312 Defaults to len(self).
313 count -- The maximum number of replacements to make. Defaults to
314 replace all occurrences.
315 bytealigned -- If True replacements will only be made on byte
316 boundaries.
317
318 Raises ValueError if old is empty or if start or end are
319 out of range.
320
321 """
322 if count == 0:
323 return 0
324 old = self._create_from_bitstype(old)
325 new = self._create_from_bitstype(new)
326 if len(old) == 0:
327 raise ValueError("Empty bitstring cannot be replaced.")
328 start, end = self._validate_slice(start, end)
329
330 if new is self:
331 # Prevent self assignment woes
332 new = copy.copy(self)
333 return self._replace(old, new, start, end, 0 if count is None else count, bytealigned)
334
335 def insert(self, bs: BitsType, pos: int) -> None:
336 """Insert bs at bit position pos.
337
338 bs -- The bitstring to insert.
339 pos -- The bit position to insert at.
340
341 Raises ValueError if pos < 0 or pos > len(self).
342
343 """
344 bs = self._create_from_bitstype(bs)
345 if len(bs) == 0:
346 return
347 if bs is self:
348 bs = self._copy()
349 if pos < 0:
350 pos += len(self)
351 if not 0 <= pos <= len(self):
352 raise ValueError("Invalid insert position.")
353 self._insert(bs, pos)
354
355 def overwrite(self, bs: BitsType, pos: int) -> None:
356 """Overwrite with bs at bit position pos.
357
358 bs -- The bitstring to overwrite with.
359 pos -- The bit position to begin overwriting from.
360
361 Raises ValueError if pos < 0 or pos > len(self).
362
363 """
364 bs = self._create_from_bitstype(bs)
365 if len(bs) == 0:
366 return
367 if pos < 0:
368 pos += len(self)
369 if pos < 0 or pos > len(self):
370 raise ValueError("Overwrite starts outside boundary of bitstring.")
371 self._overwrite(bs, pos)
372
373 def append(self, bs: BitsType) -> None:
374 """Append a bitstring to the current bitstring.
375
376 bs -- The bitstring to append.
377
378 """
379 self._append(bs)
380
381 def prepend(self, bs: BitsType) -> None:
382 """Prepend a bitstring to the current bitstring.
383
384 bs -- The bitstring to prepend.
385
386 """
387 self._prepend(bs)
388
389 def _append_msb0(self, bs: BitsType) -> None:
390 self._addright(self._create_from_bitstype(bs))
391
392 def _append_lsb0(self, bs: BitsType) -> None:
393 bs = self._create_from_bitstype(bs)
394 self._addleft(bs)
395
396 def reverse(self, start: Optional[int] = None, end: Optional[int] = None) -> None:
397 """Reverse bits in-place.
398
399 start -- Position of first bit to reverse. Defaults to 0.
400 end -- One past the position of the last bit to reverse.
401 Defaults to len(self).
402
403 Using on an empty bitstring will have no effect.
404
405 Raises ValueError if start < 0, end > len(self) or end < start.
406
407 """
408 start, end = self._validate_slice(start, end)
409 if start == 0 and end == len(self):
410 self._bitstore.reverse()
411 return
412 s = self._slice(start, end)
413 s._bitstore.reverse()
414 self[start:end] = s
415
416 def set(self, value: Any, pos: Optional[Union[int, Iterable[int]]] = None) -> None:
417 """Set one or many bits to 1 or 0.
418
419 value -- If bool(value) is True bits are set to 1, otherwise they are set to 0.
420 pos -- Either a single bit position or an iterable of bit positions.
421 Negative numbers are treated in the same way as slice indices.
422 Defaults to the entire bitstring.
423
424 Raises IndexError if pos < -len(self) or pos >= len(self).
425
426 """
427 if pos is None:
428 # Set all bits to either 1 or 0
429 self._setint(-1 if value else 0)
430 return
431 if not isinstance(pos, abc.Iterable):
432 pos = (pos,)
433 v = 1 if value else 0
434 if isinstance(pos, range):
435 self._bitstore.__setitem__(slice(pos.start, pos.stop, pos.step), v)
436 return
437 for p in pos:
438 self._bitstore[p] = v
439
440 def invert(self, pos: Optional[Union[Iterable[int], int]] = None) -> None:
441 """Invert one or many bits from 0 to 1 or vice versa.
442
443 pos -- Either a single bit position or an iterable of bit positions.
444 Negative numbers are treated in the same way as slice indices.
445
446 Raises IndexError if pos < -len(self) or pos >= len(self).
447
448 """
449 if pos is None:
450 self._invert_all()
451 return
452 if not isinstance(pos, abc.Iterable):
453 pos = (pos,)
454 length = len(self)
455
456 for p in pos:
457 if p < 0:
458 p += length
459 if not 0 <= p < length:
460 raise IndexError(f"Bit position {p} out of range.")
461 self._invert(p)
462
463 def ror(self, bits: int, start: Optional[int] = None, end: Optional[int] = None) -> None:
464 """Rotate bits to the right in-place.
465
466 bits -- The number of bits to rotate by.
467 start -- Start of slice to rotate. Defaults to 0.
468 end -- End of slice to rotate. Defaults to len(self).
469
470 Raises ValueError if bits < 0.
471
472 """
473 if not len(self):
474 raise Error("Cannot rotate an empty bitstring.")
475 if bits < 0:
476 raise ValueError("Cannot rotate by negative amount.")
477 self._ror(bits, start, end)
478
479 def _ror_msb0(self, bits: int, start: Optional[int] = None, end: Optional[int] = None) -> None:
480 start, end = self._validate_slice(start, end) # the _slice deals with msb0/lsb0
481 bits %= (end - start)
482 if not bits:
483 return
484 rhs = self._slice(end - bits, end)
485 self._delete(bits, end - bits)
486 self._insert(rhs, start)
487
488 def rol(self, bits: int, start: Optional[int] = None, end: Optional[int] = None) -> None:
489 """Rotate bits to the left in-place.
490
491 bits -- The number of bits to rotate by.
492 start -- Start of slice to rotate. Defaults to 0.
493 end -- End of slice to rotate. Defaults to len(self).
494
495 Raises ValueError if bits < 0.
496
497 """
498 if not len(self):
499 raise Error("Cannot rotate an empty bitstring.")
500 if bits < 0:
501 raise ValueError("Cannot rotate by negative amount.")
502 self._rol(bits, start, end)
503
504 def _rol_msb0(self, bits: int, start: Optional[int] = None, end: Optional[int] = None):
505 start, end = self._validate_slice(start, end)
506 bits %= (end - start)
507 if bits == 0:
508 return
509 lhs = self._slice(start, start + bits)
510 self._delete(bits, start)
511 self._insert(lhs, end - bits)
512
513 def byteswap(self, fmt: Optional[Union[int, Iterable[int], str]] = None, start: Optional[int] = None,
514 end: Optional[int] = None, repeat: bool = True) -> int:
515 """Change the endianness in-place. Return number of repeats of fmt done.
516
517 fmt -- A compact structure string, an integer number of bytes or
518 an iterable of integers. Defaults to 0, which byte reverses the
519 whole bitstring.
520 start -- Start bit position, defaults to 0.
521 end -- End bit position, defaults to len(self).
522 repeat -- If True (the default) the byte swapping pattern is repeated
523 as much as possible.
524
525 """
526 start_v, end_v = self._validate_slice(start, end)
527 if fmt is None or fmt == 0:
528 # reverse all of the whole bytes.
529 bytesizes = [(end_v - start_v) // 8]
530 elif isinstance(fmt, numbers.Integral):
531 if fmt < 0:
532 raise ValueError(f"Improper byte length {fmt}.")
533 bytesizes = [fmt]
534 elif isinstance(fmt, str):
535 if not (m := utils.BYTESWAP_STRUCT_PACK_RE.match(fmt)):
536 raise ValueError(f"Cannot parse format string {fmt}.")
537 # Split the format string into a list of 'q', '4h' etc.
538 formatlist = re.findall(utils.STRUCT_SPLIT_RE, m.group('fmt'))
539 # Now deal with multiplicative factors, 4h -> hhhh etc.
540 bytesizes = []
541 for f in formatlist:
542 if len(f) == 1:
543 bytesizes.append(utils.PACK_CODE_SIZE[f])
544 else:
545 bytesizes.extend([utils.PACK_CODE_SIZE[f[-1]]] * int(f[:-1]))
546 elif isinstance(fmt, abc.Iterable):
547 bytesizes = fmt
548 for bytesize in bytesizes:
549 if not isinstance(bytesize, numbers.Integral) or bytesize < 0:
550 raise ValueError(f"Improper byte length {bytesize}.")
551 else:
552 raise TypeError("Format must be an integer, string or iterable.")
553
554 repeats = 0
555 totalbitsize: int = 8 * sum(bytesizes)
556 if not totalbitsize:
557 return 0
558 if repeat:
559 # Try to repeat up to the end of the bitstring.
560 finalbit = end_v
561 else:
562 # Just try one (set of) byteswap(s).
563 finalbit = start_v + totalbitsize
564 for patternend in range(start_v + totalbitsize, finalbit + 1, totalbitsize):
565 bytestart = patternend - totalbitsize
566 for bytesize in bytesizes:
567 byteend = bytestart + bytesize * 8
568 self._reversebytes(bytestart, byteend)
569 bytestart += bytesize * 8
570 repeats += 1
571 return repeats
572
573 def clear(self) -> None:
574 """Remove all bits, reset to zero length."""
575 self._clear()
576
577