fix: pointer subclassing now uses correct size for arithmetic operations

Author: MrIndeciso

libdestruct/common/ptr/ptr.py

@@ -149,12 +149,12 @@ def _element_size(self: ptr) -> int:
     def __add__(self: ptr, n: int) -> ptr:
         """Return a new pointer advanced by n elements."""
         new_addr = self.get() + n * self._element_size
-        return ptr(_ArithmeticResolver(self.resolver, new_addr), self.wrapper)
+        return type(self)(_ArithmeticResolver(self.resolver, new_addr), self.wrapper)
 
     def __sub__(self: ptr, n: int) -> ptr:
         """Return a new pointer retreated by n elements."""
         new_addr = self.get() - n * self._element_size
-        return ptr(_ArithmeticResolver(self.resolver, new_addr), self.wrapper)
+        return type(self)(_ArithmeticResolver(self.resolver, new_addr), self.wrapper)
 
     def __getitem__(self: ptr, n: int) -> obj:
         """Return the object at index n relative to this pointer."""

test/scripts/types_test.py

@@ -970,5 +970,26 @@ def test_unwrap_bytes_returns_live_data(self):
         self.assertEqual(p.unwrap(4), b"BBBB")
 
 
+class PtrArithmeticSubclassTest(unittest.TestCase):
+    """Pointer arithmetic must preserve subclass identity (e.g. for narrower pointer widths)."""
+
+    def test_add_returns_same_subclass(self):
+        class ptr32(ptr):
+            size: int = 4
+
+        memory = bytearray(32)
+        p = ptr32(MemoryResolver(memory, 0), c_int)
+        self.assertIsInstance(p + 1, ptr32)
+        self.assertEqual((p + 1).size, 4)
+
+    def test_sub_returns_same_subclass(self):
+        class ptr32(ptr):
+            size: int = 4
+
+        memory = bytearray(32)
+        p = ptr32(MemoryResolver(memory, 0), c_int)
+        self.assertIsInstance(p - 1, ptr32)
+
+
 if __name__ == "__main__":
     unittest.main()