Rewrite of __setattr__ with Dtypes.

scott-griffiths · Dec 30, 2023 · 973a764 · 973a764
1 parent f2580a9
commit 973a764
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 27 deletions.
diff --git a/bitstring/bitarray_.py b/bitstring/bitarray_.py
@@ -163,26 +163,17 @@ def __setattr__(self, attribute, value) -> None:
             # First try the ordinary attribute setter
             super().__setattr__(attribute, value)
         except AttributeError:
-            if name_length := BitArray._name_length_pattern.match(attribute):
-                name = name_length.group('name')
-                length = name_length.group('len')
-                if length is not None:
-                    length = int(length)
-                    if name == 'bytes':
-                        if len(value) != length:
-                            raise CreationError(
-                                f"Wrong amount of byte data preset - {length} bytes needed, have {len(value)} bytes.")
-                        length *= 8
-                try:
-                    x = Bits(length=length, offset=None, **{name: value})
-                    if len(x) != length:
-                        raise CreationError(f"Can't initialise with value of length {len(x)} bits, "
-                                            f"as attribute has length of {length} bits.")
-                    self._bitstore = x._bitstore
-                    return
-                except AttributeError:
-                    pass
-            raise AttributeError(f"Can't set attribute {attribute} with value {value}.")
+            from bitstring.dtypes import Dtype
+            dtype = Dtype(attribute)
+            x = object.__new__(Bits)
+            if (set_fn := dtype.set_fn) is None:
+                raise AttributeError(f"Cannot set attribute '{attribute}' as it does not have a set_fn.")
+            set_fn(x, value)
+            if len(x) != dtype.bitlength:
+                raise CreationError(f"Can't initialise with value of length {len(x)} bits, "
+                                    f"as attribute has length of {dtype.bitlength} bits.")
+            self._bitstore = x._bitstore
+            return
 
     def __iadd__(self, bs: BitsType) -> BitArray:
         """Append bs to current bitstring. Return self.

diff --git a/bitstring/bits.py b/bitstring/bits.py
@@ -209,6 +209,9 @@ def _initialise(self, auto: Any, /, length: Optional[int], offset: Optional[int]
                                     f"Instead of '{self.__class__.__name__}(auto=x)' use '{self.__class__.__name__}(x)'.")
             else:
                 raise CreationError(f"Unrecognised keyword '{k}' used to initialise.")
+        if k == 'bytes':
+            # Special case for bytes as we want to allow offsets and lengths to work on creation.
+            setting_function = Bits._setbytes_with_truncation
         setting_function(self, v, length, offset)
 
     def __getattr__(self, attribute: str) -> Any:
@@ -649,12 +652,15 @@ def _sete5m2float(self, f: float, length: None = None, _offset: None = None):
     def _sete4m3float(self, f: float, length: None = None, _offset: None = None):
         self._bitstore = e4m3float2bitstore(f)
 
-    def _setbytes(self, data: Union[bytearray, bytes],
-                  length: Optional[int] = None, offset: Optional[int] = None) -> None:
+    def _setbytes(self, data: Union[bytearray, bytes], length:None = None, _offset: None = None) -> None:
         """Set the data from a bytes or bytearray object."""
+        self._bitstore = BitStore(frombytes=bytearray(data))
+        return
+
+    def _setbytes_with_truncation(self, data: Union[bytearray, bytes], length: Optional[int] = None, offset: Optional[int] = None) -> None:
+        """Set the data from a bytes or bytearray object, with optional offset and length truncations."""
         if offset is None and length is None:
-            self._bitstore = BitStore(frombytes=bytearray(data))
-            return
+            return self._setbytes(data)
         data = bytearray(data)
         if offset is None:
             offset = 0

diff --git a/tests/test_dtypes.py b/tests/test_dtypes.py
@@ -79,7 +79,8 @@ def get_fn(bs):
         md = DtypeDefinition('counter', None, get_fn)
         bs.dtype_register.add_dtype(md)
         a = bs.BitStream('0x010f')
-        c = a.counter
-        self.assertEqual(c, 5)
+        self.assertEqual(a.counter, 5)
         self.assertEqual(a.readlist('2*counter8'), [1, 4])
-        self.assertEqual(a.unpack('counter7, counter'), [0, 5])
+        self.assertEqual(a.unpack('counter7, counter'), [0, 5])
+        with self.assertRaises(AttributeError):
+            a.counter = 4