Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/bitstring/array_.py: 20%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import math
4import numbers
5from collections.abc import Sized
6from bitstring.exceptions import CreationError
7from typing import Union, List, Iterable, Any, Optional, BinaryIO, overload, TextIO
8from bitstring.bits import Bits, BitsType
9from bitstring.bitarray_ import BitArray
10from bitstring.dtypes import Dtype, dtype_register
11from bitstring import utils
12from bitstring.bitstring_options import Options, Colour
13import copy
14import array
15import operator
16import io
17import sys
19# The possible types stored in each element of the Array
20ElementType = Union[float, str, int, bytes, bool, Bits]
22options = Options()
25class Array:
26 """Return an Array whose elements are initialised according to the fmt string.
27 The dtype string can be typecode as used in the struct module or any fixed-length bitstring
28 format.
30 a = Array('>H', [1, 15, 105])
31 b = Array('int5', [-9, 0, 4])
33 The Array data is stored compactly as a BitArray object and the Array behaves very like
34 a list of items of the given format. Both the Array data and fmt properties can be freely
35 modified after creation. If the data length is not a multiple of the fmt length then the
36 Array will have 'trailing_bits' which will prevent some methods from appending to the
37 Array.
39 Methods:
41 append() -- Append a single item to the end of the Array.
42 byteswap() -- Change byte endianness of all items.
43 count() -- Count the number of occurences of a value.
44 extend() -- Append new items to the end of the Array from an iterable.
45 fromfile() -- Append items read from a file object.
46 insert() -- Insert an item at a given position.
47 pop() -- Remove and return an item.
48 pp() -- Pretty print the Array.
49 reverse() -- Reverse the order of all items.
50 tobytes() -- Return Array data as bytes object, padding with zero bits at the end if needed.
51 tofile() -- Write Array data to a file, padding with zero bits at the end if needed.
52 tolist() -- Return Array items as a list.
54 Special methods:
56 Also available are the operators [], ==, !=, +, *, <<, >>, &, |, ^,
57 plus the mutating operators [], +=, *=, <<=, >>=, &=, |=, ^=.
59 Properties:
61 data -- The BitArray binary data of the Array. Can be freely modified.
62 dtype -- The format string or typecode. Can be freely modified.
63 itemsize -- The length *in bits* of a single item. Read only.
64 trailing_bits -- If the data length is not a multiple of the fmt length, this BitArray
65 gives the leftovers at the end of the data.
68 """
70 def __init__(self, dtype: Union[str, Dtype], initializer: Optional[Union[int, Array, array.array, Iterable, Bits, bytes, bytearray, memoryview, BinaryIO]] = None,
71 trailing_bits: Optional[BitsType] = None) -> None:
72 self.data = BitArray()
73 if isinstance(dtype, Dtype) and dtype.scale == 'auto':
74 if isinstance(initializer, (int, Bits, bytes, bytearray, memoryview, BinaryIO)):
75 raise TypeError("An Array with an 'auto' scale factor can only be created from an iterable of values.")
76 auto_scale = self._calculate_auto_scale(initializer, dtype.name, dtype.length)
77 dtype = Dtype(dtype.name, dtype.length, scale=auto_scale)
78 try:
79 self._set_dtype(dtype)
80 except ValueError as e:
81 raise CreationError(e)
83 if isinstance(initializer, numbers.Integral):
84 self.data = BitArray(initializer * self._dtype.bitlength)
85 elif isinstance(initializer, (Bits, bytes, bytearray, memoryview)):
86 self.data += initializer
87 elif isinstance(initializer, io.BufferedReader):
88 self.fromfile(initializer)
89 elif initializer is not None:
90 self.extend(initializer)
92 if trailing_bits is not None:
93 self.data += BitArray._create_from_bitstype(trailing_bits)
95 _largest_values = None
96 @staticmethod
97 def _calculate_auto_scale(initializer, name: str, length: Optional[int]) -> float:
98 # Now need to find the largest power of 2 representable with this format.
99 if Array._largest_values is None:
100 Array._largest_values = {
101 'mxint8': Bits('0b01111111').mxint8, # 1.0 + 63.0/64.0,
102 'e2m1mxfp4': Bits('0b0111').e2m1mxfp4, # 6.0
103 'e2m3mxfp6': Bits('0b011111').e2m3mxfp6, # 7.5
104 'e3m2mxfp6': Bits('0b011111').e3m2mxfp6, # 28.0
105 'e4m3mxfp8': Bits('0b01111110').e4m3mxfp8, # 448.0
106 'e5m2mxfp8': Bits('0b01111011').e5m2mxfp8, # 57344.0
107 'p4binary8': Bits('0b01111110').p4binary8, # 224.0
108 'p3binary8': Bits('0b01111110').p3binary8, # 49152.0
109 'float16': Bits('0x7bff').float16, # 65504.0
110 # The bfloat range is so large the scaling algorithm doesn't work well, so I'm disallowing it.
111 # 'bfloat16': Bits('0x7f7f').bfloat16, # 3.38953139e38,
112 }
113 if f'{name}{length}' in Array._largest_values.keys():
114 float_values = Array('float64', initializer).tolist()
115 max_float_value = max(abs(x) for x in float_values)
116 if max_float_value == 0:
117 # This special case isn't covered in the standard. I'm choosing to return no scale.
118 return 1.0
119 log2 = int(math.log2(max_float_value))
120 lp2 = int(math.log2(Array._largest_values[f'{name}{length}']))
121 lg_scale = log2 - lp2
122 # Saturate at values representable in E8M0 format.
123 if lg_scale > 127:
124 lg_scale = 127
125 elif lg_scale < -127:
126 lg_scale = -127
127 return 2 ** lg_scale
128 else:
129 raise ValueError(f"Can't calculate auto scale for format '{name}{length}'. "
130 f"This feature is only available for these formats: {list(Array._largest_values.keys())}.")
132 @property
133 def itemsize(self) -> int:
134 return self._dtype.length
136 @property
137 def trailing_bits(self) -> BitArray:
138 trailing_bit_length = len(self.data) % self._dtype.bitlength
139 return BitArray() if trailing_bit_length == 0 else self.data[-trailing_bit_length:]
141 @property
142 def dtype(self) -> Dtype:
143 return self._dtype
145 @dtype.setter
146 def dtype(self, new_dtype: Union[str, Dtype]) -> None:
147 self._set_dtype(new_dtype)
149 def _set_dtype(self, new_dtype: Union[str, Dtype]) -> None:
150 if isinstance(new_dtype, Dtype):
151 self._dtype = new_dtype
152 else:
153 try:
154 dtype = Dtype(new_dtype)
155 except ValueError:
156 name_length = utils.parse_single_struct_token(new_dtype)
157 if name_length is not None:
158 dtype = Dtype(*name_length)
159 else:
160 raise ValueError(f"Inappropriate Dtype for Array: '{new_dtype}'.")
161 if dtype.length is None:
162 raise ValueError(f"A fixed length format is needed for an Array, received '{new_dtype}'.")
163 self._dtype = dtype
164 if self._dtype.scale == 'auto':
165 raise ValueError("A Dtype with an 'auto' scale factor can only be used when creating a new Array.")
167 def _create_element(self, value: ElementType) -> Bits:
168 """Create Bits from value according to the token_name and token_length"""
169 b = self._dtype.build(value)
170 if len(b) != self._dtype.length:
171 raise ValueError(f"The value {value!r} has the wrong length for the format '{self._dtype}'.")
172 return b
174 def __len__(self) -> int:
175 return len(self.data) // self._dtype.length
177 @overload
178 def __getitem__(self, key: slice) -> Array:
179 ...
181 @overload
182 def __getitem__(self, key: int) -> ElementType:
183 ...
185 def __getitem__(self, key: Union[slice, int]) -> Union[Array, ElementType]:
186 if isinstance(key, slice):
187 start, stop, step = key.indices(len(self))
188 if step != 1:
189 d = BitArray()
190 for s in range(start * self._dtype.length, stop * self._dtype.length, step * self._dtype.length):
191 d.append(self.data[s: s + self._dtype.length])
192 a = self.__class__(self._dtype)
193 a.data = d
194 return a
195 else:
196 a = self.__class__(self._dtype)
197 a.data = self.data[start * self._dtype.length: stop * self._dtype.length]
198 return a
199 else:
200 if key < 0:
201 key += len(self)
202 if key < 0 or key >= len(self):
203 raise IndexError(f"Index {key} out of range for Array of length {len(self)}.")
204 return self._dtype.read_fn(self.data, start=self._dtype.length * key)
206 @overload
207 def __setitem__(self, key: slice, value: Iterable[ElementType]) -> None:
208 ...
210 @overload
211 def __setitem__(self, key: int, value: ElementType) -> None:
212 ...
214 def __setitem__(self, key: Union[slice, int], value: Union[Iterable[ElementType], ElementType]) -> None:
215 if isinstance(key, slice):
216 start, stop, step = key.indices(len(self))
217 if not isinstance(value, Iterable):
218 raise TypeError("Can only assign an iterable to a slice.")
219 if step == 1:
220 new_data = BitArray()
221 for x in value:
222 new_data += self._create_element(x)
223 self.data[start * self._dtype.length: stop * self._dtype.length] = new_data
224 return
225 items_in_slice = len(range(start, stop, step))
226 if not isinstance(value, Sized):
227 value = list(value)
228 if len(value) == items_in_slice:
229 for s, v in zip(range(start, stop, step), value):
230 self.data.overwrite(self._create_element(v), s * self._dtype.length)
231 else:
232 raise ValueError(f"Can't assign {len(value)} values to an extended slice of length {items_in_slice}.")
233 else:
234 if key < 0:
235 key += len(self)
236 if key < 0 or key >= len(self):
237 raise IndexError(f"Index {key} out of range for Array of length {len(self)}.")
238 start = self._dtype.length * key
239 self.data.overwrite(self._create_element(value), start)
240 return
242 def __delitem__(self, key: Union[slice, int]) -> None:
243 if isinstance(key, slice):
244 start, stop, step = key.indices(len(self))
245 if step == 1:
246 self.data.__delitem__(slice(start * self._dtype.length, stop * self._dtype.length))
247 return
248 # We need to delete from the end or the earlier positions will change
249 r = reversed(range(start, stop, step)) if step > 0 else range(start, stop, step)
250 for s in r:
251 self.data.__delitem__(slice(s * self._dtype.length, (s + 1) * self._dtype.length))
252 else:
253 if key < 0:
254 key += len(self)
255 if key < 0 or key >= len(self):
256 raise IndexError
257 start = self._dtype.length * key
258 del self.data[start: start + self._dtype.length]
260 def __repr__(self) -> str:
261 list_str = f"{self.tolist()}"
262 trailing_bit_length = len(self.data) % self._dtype.length
263 final_str = "" if trailing_bit_length == 0 else ", trailing_bits=" + repr(
264 self.data[-trailing_bit_length:])
265 return f"Array('{self._dtype}', {list_str}{final_str})"
267 def astype(self, dtype: Union[str, Dtype]) -> Array:
268 """Return Array with elements of new dtype, initialised from current Array."""
269 new_array = self.__class__(dtype, self.tolist())
270 return new_array
272 def tolist(self) -> List[ElementType]:
273 return [self._dtype.read_fn(self.data, start=start)
274 for start in range(0, len(self.data) - self._dtype.length + 1, self._dtype.length)]
276 def append(self, x: ElementType) -> None:
277 if len(self.data) % self._dtype.length != 0:
278 raise ValueError("Cannot append to Array as its length is not a multiple of the format length.")
279 self.data += self._create_element(x)
281 def extend(self, iterable: Union[Array, array.array, Iterable[Any]]) -> None:
282 if len(self.data) % self._dtype.length != 0:
283 raise ValueError(f"Cannot extend Array as its data length ({len(self.data)} bits) is not a multiple of the format length ({self._dtype.length} bits).")
284 if isinstance(iterable, Array):
285 if self._dtype.name != iterable._dtype.name or self._dtype.length != iterable._dtype.length:
286 raise TypeError(
287 f"Cannot extend an Array with format '{self._dtype}' from an Array of format '{iterable._dtype}'.")
288 # No need to iterate over the elements, we can just append the data
289 self.data.append(iterable.data)
290 elif isinstance(iterable, array.array):
291 # array.array types are always native-endian, hence the '='
292 name_value = utils.parse_single_struct_token('=' + iterable.typecode)
293 if name_value is None:
294 raise ValueError(f"Cannot extend from array with typecode {iterable.typecode}.")
295 other_dtype = dtype_register.get_dtype(*name_value, scale=None)
296 if self._dtype.name != other_dtype.name or self._dtype.length != other_dtype.length:
297 raise ValueError(
298 f"Cannot extend an Array with format '{self._dtype}' from an array with typecode '{iterable.typecode}'.")
299 self.data += iterable.tobytes()
300 else:
301 if isinstance(iterable, str):
302 raise TypeError("Can't extend an Array with a str.")
303 for item in iterable:
304 self.data += self._create_element(item)
306 def insert(self, i: int, x: ElementType) -> None:
307 """Insert a new element into the Array at position i.
309 """
310 i = min(i, len(self)) # Inserting beyond len of array inserts at the end (copying standard behaviour)
311 self.data.insert(self._create_element(x), i * self._dtype.length)
313 def pop(self, i: int = -1) -> ElementType:
314 """Return and remove an element of the Array.
316 Default is to return and remove the final element.
318 """
319 if len(self) == 0:
320 raise IndexError("Can't pop from an empty Array.")
321 x = self[i]
322 del self[i]
323 return x
325 def byteswap(self) -> None:
326 """Change the endianness in-place of all items in the Array.
328 If the Array format is not a whole number of bytes a ValueError will be raised.
330 """
331 if self._dtype.length % 8 != 0:
332 raise ValueError(
333 f"byteswap can only be used for whole-byte elements. The '{self._dtype}' format is {self._dtype.length} bits long.")
334 self.data.byteswap(self.itemsize // 8)
336 def count(self, value: ElementType) -> int:
337 """Return count of Array items that equal value.
339 value -- The quantity to compare each Array element to. Type should be appropriate for the Array format.
341 For floating point types using a value of float('nan') will count the number of elements that are NaN.
343 """
344 if math.isnan(value):
345 return sum(math.isnan(i) for i in self)
346 else:
347 return sum(i == value for i in self)
349 def tobytes(self) -> bytes:
350 """Return the Array data as a bytes object, padding with zero bits if needed.
352 Up to seven zero bits will be added at the end to byte align.
354 """
355 return self.data.tobytes()
357 def tofile(self, f: BinaryIO) -> None:
358 """Write the Array data to a file object, padding with zero bits if needed.
360 Up to seven zero bits will be added at the end to byte align.
362 """
363 self.data.tofile(f)
365 def fromfile(self, f: BinaryIO, n: Optional[int] = None) -> None:
366 trailing_bit_length = len(self.data) % self._dtype.bitlength
367 if trailing_bit_length != 0:
368 raise ValueError(f"Cannot extend Array as its data length ({len(self.data)} bits) is not a multiple of the format length ({self._dtype.bitlength} bits).")
370 new_data = Bits(f)
371 max_items = len(new_data) // self._dtype.length
372 items_to_append = max_items if n is None else min(n, max_items)
373 self.data += new_data[0: items_to_append * self._dtype.bitlength]
374 if n is not None and items_to_append < n:
375 raise EOFError(f"Only {items_to_append} were appended, not the {n} items requested.")
377 def reverse(self) -> None:
378 trailing_bit_length = len(self.data) % self._dtype.length
379 if trailing_bit_length != 0:
380 raise ValueError(f"Cannot reverse the items in the Array as its data length ({len(self.data)} bits) is not a multiple of the format length ({self._dtype.length} bits).")
381 for start_bit in range(0, len(self.data) // 2, self._dtype.length):
382 start_swap_bit = len(self.data) - start_bit - self._dtype.length
383 temp = self.data[start_bit: start_bit + self._dtype.length]
384 self.data[start_bit: start_bit + self._dtype.length] = self.data[
385 start_swap_bit: start_swap_bit + self._dtype.length]
386 self.data[start_swap_bit: start_swap_bit + self._dtype.length] = temp
388 def pp(self, fmt: Optional[str] = None, width: int = 120,
389 show_offset: bool = True, stream: TextIO = sys.stdout) -> None:
390 """Pretty-print the Array contents.
392 fmt -- Data format string. Defaults to current Array dtype.
393 width -- Max width of printed lines in characters. Defaults to 120. A single group will always
394 be printed per line even if it exceeds the max width.
395 show_offset -- If True shows the element offset in the first column of each line.
396 stream -- A TextIO object with a write() method. Defaults to sys.stdout.
398 """
399 colour = Colour(not options.no_color)
400 sep = ' '
401 dtype2 = None
402 tidy_fmt = None
403 if fmt is None:
404 fmt = self.dtype
405 dtype1 = self.dtype
406 tidy_fmt = "dtype='" + colour.purple + str(self.dtype) + "'" + colour.off
407 else:
408 token_list = utils.preprocess_tokens(fmt)
409 if len(token_list) not in [1, 2]:
410 raise ValueError(f"Only one or two tokens can be used in an Array.pp() format - '{fmt}' has {len(token_list)} tokens.")
411 dtype1 = Dtype(*utils.parse_name_length_token(token_list[0]))
412 if len(token_list) == 2:
413 dtype2 = Dtype(*utils.parse_name_length_token(token_list[1]))
415 token_length = dtype1.bitlength
416 if dtype2 is not None:
417 # For two types we're OK as long as they don't have different lengths given.
418 if dtype1.bitlength is not None and dtype2.bitlength is not None and dtype1.bitlength != dtype2.bitlength:
419 raise ValueError(f"Two different format lengths specified ('{fmt}'). Either specify just one, or two the same length.")
420 if token_length is None:
421 token_length = dtype2.bitlength
422 if token_length is None:
423 token_length = self.itemsize
425 trailing_bit_length = len(self.data) % token_length
426 format_sep = " : " # String to insert on each line between multiple formats
427 if tidy_fmt is None:
428 tidy_fmt = colour.purple + str(dtype1) + colour.off
429 if dtype2 is not None:
430 tidy_fmt += ', ' + colour.blue + str(dtype2) + colour.off
431 tidy_fmt = "fmt='" + tidy_fmt + "'"
432 data = self.data if trailing_bit_length == 0 else self.data[0: -trailing_bit_length]
433 length = len(self.data) // token_length
434 len_str = colour.green + str(length) + colour.off
435 stream.write(f"<{self.__class__.__name__} {tidy_fmt}, length={len_str}, itemsize={token_length} bits, total data size={(len(self.data) + 7) // 8} bytes> [\n")
436 data._pp(dtype1, dtype2, token_length, width, sep, format_sep, show_offset, stream, False, token_length)
437 stream.write("]")
438 if trailing_bit_length != 0:
439 stream.write(" + trailing_bits = " + str(self.data[-trailing_bit_length:]))
440 stream.write("\n")
442 def equals(self, other: Any) -> bool:
443 """Return True if format and all Array items are equal."""
444 if isinstance(other, Array):
445 if self._dtype.length != other._dtype.length:
446 return False
447 if self._dtype.name != other._dtype.name:
448 return False
449 if self.data != other.data:
450 return False
451 return True
452 elif isinstance(other, array.array):
453 # Assume we are comparing with an array type
454 if self.trailing_bits:
455 return False
456 # array's itemsize is in bytes, not bits.
457 if self.itemsize != other.itemsize * 8:
458 return False
459 if len(self) != len(other):
460 return False
461 if self.tolist() != other.tolist():
462 return False
463 return True
464 return False
466 def __iter__(self) -> Iterable[ElementType]:
467 start = 0
468 for _ in range(len(self)):
469 yield self._dtype.read_fn(self.data, start=start)
470 start += self._dtype.length
472 def __copy__(self) -> Array:
473 a_copy = self.__class__(self._dtype)
474 a_copy.data = copy.copy(self.data)
475 return a_copy
477 def _apply_op_to_all_elements(self, op, value: Union[int, float, None], is_comparison: bool = False) -> Array:
478 """Apply op with value to each element of the Array and return a new Array"""
479 new_array = self.__class__('bool' if is_comparison else self._dtype)
480 new_data = BitArray()
481 failures = index = 0
482 msg = ''
483 if value is not None:
484 def partial_op(a):
485 return op(a, value)
486 else:
487 def partial_op(a):
488 return op(a)
489 for i in range(len(self)):
490 v = self._dtype.read_fn(self.data, start=self._dtype.length * i)
491 try:
492 new_data.append(new_array._create_element(partial_op(v)))
493 except (CreationError, ZeroDivisionError, ValueError) as e:
494 if failures == 0:
495 msg = str(e)
496 index = i
497 failures += 1
498 if failures != 0:
499 raise ValueError(f"Applying operator '{op.__name__}' to Array caused {failures} errors. "
500 f'First error at index {index} was: "{msg}"')
501 new_array.data = new_data
502 return new_array
504 def _apply_op_to_all_elements_inplace(self, op, value: Union[int, float]) -> Array:
505 """Apply op with value to each element of the Array in place."""
506 # This isn't really being done in-place, but it's simpler and faster for now?
507 new_data = BitArray()
508 failures = index = 0
509 msg = ''
510 for i in range(len(self)):
511 v = self._dtype.read_fn(self.data, start=self._dtype.length * i)
512 try:
513 new_data.append(self._create_element(op(v, value)))
514 except (CreationError, ZeroDivisionError, ValueError) as e:
515 if failures == 0:
516 msg = str(e)
517 index = i
518 failures += 1
519 if failures != 0:
520 raise ValueError(f"Applying operator '{op.__name__}' to Array caused {failures} errors. "
521 f'First error at index {index} was: "{msg}"')
522 self.data = new_data
523 return self
525 def _apply_bitwise_op_to_all_elements(self, op, value: BitsType) -> Array:
526 """Apply op with value to each element of the Array as an unsigned integer and return a new Array"""
527 a_copy = self[:]
528 a_copy._apply_bitwise_op_to_all_elements_inplace(op, value)
529 return a_copy
531 def _apply_bitwise_op_to_all_elements_inplace(self, op, value: BitsType) -> Array:
532 """Apply op with value to each element of the Array as an unsigned integer in place."""
533 value = BitArray._create_from_bitstype(value)
534 if len(value) != self._dtype.length:
535 raise ValueError(f"Bitwise op needs a bitstring of length {self._dtype.length} to match format {self._dtype}.")
536 for start in range(0, len(self) * self._dtype.length, self._dtype.length):
537 self.data[start: start + self._dtype.length] = op(self.data[start: start + self._dtype.length], value)
538 return self
540 def _apply_op_between_arrays(self, op, other: Array, is_comparison: bool = False) -> Array:
541 if len(self) != len(other):
542 msg = f"Cannot operate element-wise on Arrays with different lengths ({len(self)} and {len(other)})."
543 if op in [operator.add, operator.iadd]:
544 msg += " Use extend() method to concatenate Arrays."
545 if op in [operator.eq, operator.ne]:
546 msg += " Use equals() method to compare Arrays for a single boolean result."
547 raise ValueError(msg)
548 if is_comparison:
549 new_type = dtype_register.get_dtype('bool', 1)
550 else:
551 new_type = self._promotetype(self._dtype, other._dtype)
552 new_array = self.__class__(new_type)
553 new_data = BitArray()
554 failures = index = 0
555 msg = ''
556 for i in range(len(self)):
557 a = self._dtype.read_fn(self.data, start=self._dtype.length * i)
558 b = other._dtype.read_fn(other.data, start=other._dtype.length * i)
559 try:
560 new_data.append(new_array._create_element(op(a, b)))
561 except (CreationError, ValueError, ZeroDivisionError) as e:
562 if failures == 0:
563 msg = str(e)
564 index = i
565 failures += 1
566 if failures != 0:
567 raise ValueError(f"Applying operator '{op.__name__}' between Arrays caused {failures} errors. "
568 f'First error at index {index} was: "{msg}"')
569 new_array.data = new_data
570 return new_array
572 @classmethod
573 def _promotetype(cls, type1: Dtype, type2: Dtype) -> Dtype:
574 """When combining types which one wins?
576 1. We only deal with types representing floats or integers.
577 2. One of the two types gets returned. We never create a new one.
578 3. Floating point types always win against integer types.
579 4. Signed integer types always win against unsigned integer types.
580 5. Longer types win against shorter types.
581 6. In a tie the first type wins against the second type.
583 """
584 def is_float(x): return x.return_type is float
585 def is_int(x): return x.return_type is int or x.return_type is bool
586 if is_float(type1) + is_int(type1) + is_float(type2) + is_int(type2) != 2:
587 raise ValueError(f"Only integer and floating point types can be combined - not '{type1}' and '{type2}'.")
588 # If same type choose the widest
589 if type1.name == type2.name:
590 return type1 if type1.length > type2.length else type2
591 # We choose floats above integers, irrespective of the widths
592 if is_float(type1) and is_int(type2):
593 return type1
594 if is_int(type1) and is_float(type2):
595 return type2
596 if is_float(type1) and is_float(type2):
597 return type2 if type2.length > type1.length else type1
598 assert is_int(type1) and is_int(type2)
599 if type1.is_signed and not type2.is_signed:
600 return type1
601 if type2.is_signed and not type1.is_signed:
602 return type2
603 return type2 if type2.length > type1.length else type1
605 # Operators between Arrays or an Array and scalar value
607 def __add__(self, other: Union[int, float, Array]) -> Array:
608 """Add int or float to all elements."""
609 if isinstance(other, Array):
610 return self._apply_op_between_arrays(operator.add, other)
611 return self._apply_op_to_all_elements(operator.add, other)
613 def __iadd__(self, other: Union[int, float, Array]) -> Array:
614 if isinstance(other, Array):
615 return self._apply_op_between_arrays(operator.add, other)
616 return self._apply_op_to_all_elements_inplace(operator.add, other)
618 def __isub__(self, other: Union[int, float, Array]) -> Array:
619 if isinstance(other, Array):
620 return self._apply_op_between_arrays(operator.sub, other)
621 return self._apply_op_to_all_elements_inplace(operator.sub, other)
623 def __sub__(self, other: Union[int, float, Array]) -> Array:
624 if isinstance(other, Array):
625 return self._apply_op_between_arrays(operator.sub, other)
626 return self._apply_op_to_all_elements(operator.sub, other)
628 def __mul__(self, other: Union[int, float, Array]) -> Array:
629 if isinstance(other, Array):
630 return self._apply_op_between_arrays(operator.mul, other)
631 return self._apply_op_to_all_elements(operator.mul, other)
633 def __imul__(self, other: Union[int, float, Array]) -> Array:
634 if isinstance(other, Array):
635 return self._apply_op_between_arrays(operator.mul, other)
636 return self._apply_op_to_all_elements_inplace(operator.mul, other)
638 def __floordiv__(self, other: Union[int, float, Array]) -> Array:
639 if isinstance(other, Array):
640 return self._apply_op_between_arrays(operator.floordiv, other)
641 return self._apply_op_to_all_elements(operator.floordiv, other)
643 def __ifloordiv__(self, other: Union[int, float, Array]) -> Array:
644 if isinstance(other, Array):
645 return self._apply_op_between_arrays(operator.floordiv, other)
646 return self._apply_op_to_all_elements_inplace(operator.floordiv, other)
648 def __truediv__(self, other: Union[int, float, Array]) -> Array:
649 if isinstance(other, Array):
650 return self._apply_op_between_arrays(operator.truediv, other)
651 return self._apply_op_to_all_elements(operator.truediv, other)
653 def __itruediv__(self, other: Union[int, float, Array]) -> Array:
654 if isinstance(other, Array):
655 return self._apply_op_between_arrays(operator.truediv, other)
656 return self._apply_op_to_all_elements_inplace(operator.truediv, other)
658 def __rshift__(self, other: Union[int, Array]) -> Array:
659 if isinstance(other, Array):
660 return self._apply_op_between_arrays(operator.rshift, other)
661 return self._apply_op_to_all_elements(operator.rshift, other)
663 def __lshift__(self, other: Union[int, Array]) -> Array:
664 if isinstance(other, Array):
665 return self._apply_op_between_arrays(operator.lshift, other)
666 return self._apply_op_to_all_elements(operator.lshift, other)
668 def __irshift__(self, other: Union[int, Array]) -> Array:
669 if isinstance(other, Array):
670 return self._apply_op_between_arrays(operator.rshift, other)
671 return self._apply_op_to_all_elements_inplace(operator.rshift, other)
673 def __ilshift__(self, other: Union[int, Array]) -> Array:
674 if isinstance(other, Array):
675 return self._apply_op_between_arrays(operator.lshift, other)
676 return self._apply_op_to_all_elements_inplace(operator.lshift, other)
678 def __mod__(self, other: Union[int, Array]) -> Array:
679 if isinstance(other, Array):
680 return self._apply_op_between_arrays(operator.mod, other)
681 return self._apply_op_to_all_elements(operator.mod, other)
683 def __imod__(self, other: Union[int, Array]) -> Array:
684 if isinstance(other, Array):
685 return self._apply_op_between_arrays(operator.mod, other)
686 return self._apply_op_to_all_elements_inplace(operator.mod, other)
688 # Bitwise operators
690 def __and__(self, other: BitsType) -> Array:
691 return self._apply_bitwise_op_to_all_elements(operator.iand, other)
693 def __iand__(self, other: BitsType) -> Array:
694 return self._apply_bitwise_op_to_all_elements_inplace(operator.iand, other)
696 def __or__(self, other: BitsType) -> Array:
697 return self._apply_bitwise_op_to_all_elements(operator.ior, other)
699 def __ior__(self, other: BitsType) -> Array:
700 return self._apply_bitwise_op_to_all_elements_inplace(operator.ior, other)
702 def __xor__(self, other: BitsType) -> Array:
703 return self._apply_bitwise_op_to_all_elements(operator.ixor, other)
705 def __ixor__(self, other: BitsType) -> Array:
706 return self._apply_bitwise_op_to_all_elements_inplace(operator.ixor, other)
708 # Reverse operators between a scalar value and an Array
710 def __rmul__(self, other: Union[int, float]) -> Array:
711 return self._apply_op_to_all_elements(operator.mul, other)
713 def __radd__(self, other: Union[int, float]) -> Array:
714 return self._apply_op_to_all_elements(operator.add, other)
716 def __rsub__(self, other: Union[int, float]) -> Array:
717 # i - A == (-A) + i
718 neg = self._apply_op_to_all_elements(operator.neg, None)
719 return neg._apply_op_to_all_elements(operator.add, other)
721 # Reverse operators between a scalar and something that can be a BitArray.
723 def __rand__(self, other: BitsType) -> Array:
724 return self._apply_bitwise_op_to_all_elements(operator.iand, other)
726 def __ror__(self, other: BitsType) -> Array:
727 return self._apply_bitwise_op_to_all_elements(operator.ior, other)
729 def __rxor__(self, other: BitsType) -> Array:
730 return self._apply_bitwise_op_to_all_elements(operator.ixor, other)
732 # Comparison operators
734 def __lt__(self, other: Union[int, float, Array]) -> Array:
735 if isinstance(other, Array):
736 return self._apply_op_between_arrays(operator.lt, other, is_comparison=True)
737 return self._apply_op_to_all_elements(operator.lt, other, is_comparison=True)
739 def __gt__(self, other: Union[int, float, Array]) -> Array:
740 if isinstance(other, Array):
741 return self._apply_op_between_arrays(operator.gt, other, is_comparison=True)
742 return self._apply_op_to_all_elements(operator.gt, other, is_comparison=True)
744 def __ge__(self, other: Union[int, float, Array]) -> Array:
745 if isinstance(other, Array):
746 return self._apply_op_between_arrays(operator.ge, other, is_comparison=True)
747 return self._apply_op_to_all_elements(operator.ge, other, is_comparison=True)
749 def __le__(self, other: Union[int, float, Array]) -> Array:
750 if isinstance(other, Array):
751 return self._apply_op_between_arrays(operator.le, other, is_comparison=True)
752 return self._apply_op_to_all_elements(operator.le, other, is_comparison=True)
754 def _eq_ne(self, op, other: Any) -> Array:
755 if isinstance(other, (int, float, str, Bits)):
756 return self._apply_op_to_all_elements(op, other, is_comparison=True)
757 try:
758 other = self.__class__(self.dtype, other)
759 except:
760 return NotImplemented
761 finally:
762 return self._apply_op_between_arrays(op, other, is_comparison=True)
764 def __eq__(self, other: Any) -> Array:
765 return self._eq_ne(operator.eq, other)
767 def __ne__(self, other: Any) -> Array:
768 return self._eq_ne(operator.ne, other)
770 # Unary operators
772 def __neg__(self):
773 return self._apply_op_to_all_elements(operator.neg, None)
775 def __abs__(self):
776 return self._apply_op_to_all_elements(operator.abs, None)