Merge pull request #7071 from youknowone/memoryview

Update test_memoryview from v3.14.3 and impl count/index

Lib/test/test_memoryview.py

@@ -13,8 +13,15 @@ import array
 import io
 import copy
 import pickle
+import struct
 
-from test.support import import_helper
+from itertools import product
+from test import support
+from test.support import import_helper, threading_helper
+
+
+class MyObject:
+    pass
 
 
 class AbstractMemoryTests:
@@ -53,12 +60,53 @@ class AbstractMemoryTests:
         for tp in self._types:
             self.check_getitem_with_type(tp)
 
+    def test_index(self):
+        for tp in self._types:
+            b = tp(self._source)
+            m = self._view(b)  # may be a sub-view
+            l = m.tolist()
+            k = 2 * len(self._source)
+
+            for chi in self._source:
+                if chi in l:
+                    self.assertEqual(m.index(chi), l.index(chi))
+                else:
+                    self.assertRaises(ValueError, m.index, chi)
+
+                for start, stop in product(range(-k, k), range(-k, k)):
+                    index = -1
+                    try:
+                        index = l.index(chi, start, stop)
+                    except ValueError:
+                        pass
+
+                    if index == -1:
+                        self.assertRaises(ValueError, m.index, chi, start, stop)
+                    else:
+                        self.assertEqual(m.index(chi, start, stop), index)
+
     def test_iter(self):
         for tp in self._types:
             b = tp(self._source)
             m = self._view(b)
             self.assertEqual(list(m), [m[i] for i in range(len(m))])
 
+    def test_count(self):
+        for tp in self._types:
+            b = tp(self._source)
+            m = self._view(b)
+            l = m.tolist()
+            for ch in list(m):
+                self.assertEqual(m.count(ch), l.count(ch))
+
+            b = tp((b'a' * 5) + (b'c' * 3))
+            m = self._view(b)  # may be sliced
+            l = m.tolist()
+            with self.subTest('count', buffer=b):
+                self.assertEqual(m.count(ord('a')), l.count(ord('a')))
+                self.assertEqual(m.count(ord('b')), l.count(ord('b')))
+                self.assertEqual(m.count(ord('c')), l.count(ord('c')))
+
     def test_setitem_readonly(self):
         if not self.ro_type:
             self.skipTest("no read-only type to test")
@@ -191,16 +239,12 @@ class AbstractMemoryTests:
         self.assertEqual(m.suboffsets, ())
         return m
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_attributes_readonly(self):
         if not self.ro_type:
             self.skipTest("no read-only type to test")
         m = self.check_attributes_with_type(self.ro_type)
         self.assertEqual(m.readonly, True)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_attributes_writable(self):
         if not self.rw_type:
             self.skipTest("no writable type to test")
@@ -231,8 +275,6 @@ class AbstractMemoryTests:
                     self.m = memoryview(base)
             class MySource(tp):
                 pass
-            class MyObject:
-                pass
 
             # Create a reference cycle through a memoryview object.
             # This exercises mbuf_clear().
@@ -345,6 +387,21 @@ class AbstractMemoryTests:
         m = self._view(b)
         self.assertRaises(ValueError, hash, m)
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; re-entrant buffer release not detected
+    def test_hash_use_after_free(self):
+        # Prevent crash in memoryview(v).__hash__ with re-entrant v.__hash__.
+        # Regression test for https://github.com/python/cpython/issues/142664.
+        class E(array.array):
+            def __hash__(self):
+                mv.release()
+                self.clear()
+                return 123
+
+        v = E('B', b'A' * 4096)
+        mv = memoryview(v).toreadonly()   # must be read-only for hash()
+        self.assertRaises(BufferError, hash, mv)
+        self.assertRaises(BufferError, mv.__hash__)
+
     def test_weakref(self):
         # Check memoryviews are weakrefable
         for tp in self._types:
@@ -400,6 +457,21 @@ class AbstractMemoryTests:
         self.assertEqual(c.format, "H")
         self.assertEqual(d.format, "H")
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; re-entrant buffer release not detected
+    def test_hex_use_after_free(self):
+        # Prevent UAF in memoryview.hex(sep) with re-entrant sep.__len__.
+        # Regression test for https://github.com/python/cpython/issues/143195.
+        ba = bytearray(b'A' * 1024)
+        mv = memoryview(ba)
+
+        class S(bytes):
+            def __len__(self):
+                mv.release()
+                ba.clear()
+                return 1
+
+        self.assertRaises(BufferError, mv.hex, S(b':'))
+
 
 # Variations on source objects for the buffer: bytes-like objects, then arrays
 # with itemsize > 1.
@@ -439,6 +511,18 @@ class BaseMemoryviewTests:
     def _check_contents(self, tp, obj, contents):
         self.assertEqual(obj, tp(contents))
 
+    def test_count(self):
+        super().test_count()
+        for tp in self._types:
+            b = tp((b'a' * 5) + (b'c' * 3))
+            m = self._view(b)  # should not be sliced
+            self.assertEqual(len(b), len(m))
+            with self.subTest('count', buffer=b):
+                self.assertEqual(m.count(ord('a')), 5)
+                self.assertEqual(m.count(ord('b')), 0)
+                self.assertEqual(m.count(ord('c')), 3)
+
+
 class BaseMemorySliceTests:
     source_bytes = b"XabcdefY"
 
@@ -472,11 +556,6 @@ class BaseMemorySliceSliceTests:
 class BytesMemoryviewTest(unittest.TestCase,
     BaseMemoryviewTests, BaseBytesMemoryTests):
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
-    def test_gc(self):
-        super().test_gc()
-
     def test_constructor(self):
         for tp in self._types:
             ob = tp(self._source)
@@ -487,6 +566,10 @@ class BytesMemoryviewTest(unittest.TestCase,
             self.assertRaises(TypeError, memoryview, argument=ob)
             self.assertRaises(TypeError, memoryview, ob, argument=True)
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: False is not true : <test.test_memoryview.MyObject object at 0x84ecb1920>
+    def test_gc(self):
+        return super().test_gc()
+
 class ArrayMemoryviewTest(unittest.TestCase,
     BaseMemoryviewTests, BaseArrayMemoryTests):
 
@@ -501,24 +584,24 @@ class ArrayMemoryviewTest(unittest.TestCase,
 
 class BytesMemorySliceTest(unittest.TestCase,
     BaseMemorySliceTests, BaseBytesMemoryTests):
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
-    def test_gc(self):
-        super().test_gc()
     pass
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: False is not true : <test.test_memoryview.MyObject object at 0x84ecb13e0>
+    def test_gc(self):
+        return super().test_gc()
+
 class ArrayMemorySliceTest(unittest.TestCase,
     BaseMemorySliceTests, BaseArrayMemoryTests):
     pass
 
 class BytesMemorySliceSliceTest(unittest.TestCase,
     BaseMemorySliceSliceTests, BaseBytesMemoryTests):
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
-    def test_gc(self):
-        super().test_gc()
     pass
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: False is not true : <test.test_memoryview.MyObject object at 0x84ddca1c0>
+    def test_gc(self):
+        return super().test_gc()
+
 class ArrayMemorySliceSliceTest(unittest.TestCase,
     BaseMemorySliceSliceTests, BaseArrayMemoryTests):
     pass
@@ -544,6 +627,14 @@ class OtherTest(unittest.TestCase):
                 m[2:] = memoryview(p6).cast(format)[2:]
                 self.assertEqual(d.value, 0.6)
 
+    def test_half_float(self):
+        half_data = struct.pack('eee', 0.0, -1.5, 1.5)
+        float_data = struct.pack('fff', 0.0, -1.5, 1.5)
+        half_view = memoryview(half_data).cast('e')
+        float_view = memoryview(float_data).cast('f')
+        self.assertEqual(half_view.nbytes * 2, float_view.nbytes)
+        self.assertListEqual(half_view.tolist(), float_view.tolist())
+
     def test_memoryview_hex(self):
         # Issue #9951: memoryview.hex() segfaults with non-contiguous buffers.
         x = b'0' * 200000
@@ -551,6 +642,26 @@ class OtherTest(unittest.TestCase):
         m2 = m1[::-1]
         self.assertEqual(m2.hex(), '30' * 200000)
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; TypeError: Unexpected keyword argument sep
+    def test_memoryview_hex_separator(self):
+        x = bytes(range(97, 102))
+        m1 = memoryview(x)
+        m2 = m1[::-1]
+        self.assertEqual(m2.hex(':'), '65:64:63:62:61')
+        self.assertEqual(m2.hex(':', 2), '65:6463:6261')
+        self.assertEqual(m2.hex(':', -2), '6564:6362:61')
+        self.assertEqual(m2.hex(sep=':', bytes_per_sep=2), '65:6463:6261')
+        self.assertEqual(m2.hex(sep=':', bytes_per_sep=-2), '6564:6362:61')
+        for bytes_per_sep in 5, -5, 2**31-1, -(2**31-1):
+            with self.subTest(bytes_per_sep=bytes_per_sep):
+                self.assertEqual(m2.hex(':', bytes_per_sep), '6564636261')
+        for bytes_per_sep in 2**31, -2**31, 2**1000, -2**1000:
+            with self.subTest(bytes_per_sep=bytes_per_sep):
+                try:
+                    self.assertEqual(m2.hex(':', bytes_per_sep), '6564636261')
+                except OverflowError:
+                    pass
+
     def test_copy(self):
         m = memoryview(b'abc')
         with self.assertRaises(TypeError):
@@ -562,8 +673,7 @@ class OtherTest(unittest.TestCase):
             with self.assertRaises(TypeError):
                 pickle.dumps(m, proto)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: ValueError not raised
     def test_use_released_memory(self):
         # gh-92888: Previously it was possible to use a memoryview even after
         # backing buffer is freed in certain cases. This tests that those
@@ -666,5 +776,56 @@ class OtherTest(unittest.TestCase):
             m[0] = MyBool()
         self.assertEqual(ba[:8], b'\0'*8)
 
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AttributeError: 'memoryview' object has no attribute '__buffer__'
+    def test_buffer_reference_loop(self):
+        m = memoryview(b'abc').__buffer__(0)
+        o = MyObject()
+        o.m = m
+        o.o = o
+        wr = weakref.ref(o)
+        del m, o
+        gc.collect()
+        self.assertIsNone(wr())
+
+    @unittest.expectedFailure  # TODO: RUSTPYTHON; AttributeError: module 'pickle' has no attribute 'PickleBuffer'
+    def test_picklebuffer_reference_loop(self):
+        pb = pickle.PickleBuffer(memoryview(b'abc'))
+        o = MyObject()
+        o.pb = pb
+        o.o = o
+        wr = weakref.ref(o)
+        del pb, o
+        gc.collect()
+        self.assertIsNone(wr())
+
+
+@threading_helper.requires_working_threading()
+@support.requires_resource("cpu")
+class RacingTest(unittest.TestCase):
+    def test_racing_getbuf_and_releasebuf(self):
+        """Repeatly access the memoryview for racing."""
+        try:
+            from multiprocessing.managers import SharedMemoryManager
+        except ImportError:
+            self.skipTest("Test requires multiprocessing")
+        from threading import Thread, Event
+
+        start = Event()
+        with SharedMemoryManager() as smm:
+            obj = smm.ShareableList(range(100))
+            def test():
+                # Issue gh-127085, the `ShareableList.count` is just a
+                # convenient way to mess the `exports` counter of `memoryview`,
+                # this issue has no direct relation with `ShareableList`.
+                start.wait(support.SHORT_TIMEOUT)
+                for i in range(10):
+                    obj.count(1)
+            threads = [Thread(target=test) for _ in range(10)]
+            with threading_helper.start_threads(threads):
+                start.set()
+
+            del obj
+
+
 if __name__ == "__main__":
     unittest.main()

crates/vm/src/builtins/memory.rs

@@ -617,13 +617,22 @@ impl PyMemoryView {
     #[pygetset]
     fn suboffsets(&self, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
         self.try_not_released(vm)?;
-        Ok(vm.ctx.new_tuple(
-            self.desc
-                .dim_desc
-                .iter()
-                .map(|(_, _, suboffset)| suboffset.to_pyobject(vm))
-                .collect(),
-        ))
+        let has_suboffsets = self
+            .desc
+            .dim_desc
+            .iter()
+            .any(|(_, _, suboffset)| *suboffset != 0);
+        if has_suboffsets {
+            Ok(vm.ctx.new_tuple(
+                self.desc
+                    .dim_desc
+                    .iter()
+                    .map(|(_, _, suboffset)| suboffset.to_pyobject(vm))
+                    .collect(),
+            ))
+        } else {
+            Ok(vm.ctx.empty_tuple.clone())
+        }
     }
 
     #[pygetset]
@@ -740,6 +749,63 @@ impl PyMemoryView {
         self.contiguous_or_collect(|x| bytes_to_hex(x, sep, bytes_per_sep, vm))
     }
 
+    #[pymethod]
+    fn count(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<usize> {
+        self.try_not_released(vm)?;
+        if self.desc.ndim() != 1 {
+            return Err(
+                vm.new_not_implemented_error("multi-dimensional sub-views are not implemented")
+            );
+        }
+        let len = self.desc.dim_desc[0].0;
+        let mut count = 0;
+        for i in 0..len {
+            let item = self.getitem_by_idx(i as isize, vm)?;
+            if vm.bool_eq(&item, &value)? {
+                count += 1;
+            }
+        }
+        Ok(count)
+    }
+
+    #[pymethod]
+    fn index(
+        &self,
+        value: PyObjectRef,
+        start: OptionalArg<isize>,
+        stop: OptionalArg<isize>,
+        vm: &VirtualMachine,
+    ) -> PyResult<usize> {
+        self.try_not_released(vm)?;
+        if self.desc.ndim() != 1 {
+            return Err(
+                vm.new_not_implemented_error("multi-dimensional sub-views are not implemented")
+            );
+        }
+        let len = self.desc.dim_desc[0].0;
+        let start = start.unwrap_or(0);
+        let stop = stop.unwrap_or(len as isize);
+
+        let start = if start < 0 {
+            (start + len as isize).max(0) as usize
+        } else {
+            (start as usize).min(len)
+        };
+        let stop = if stop < 0 {
+            (stop + len as isize).max(0) as usize
+        } else {
+            (stop as usize).min(len)
+        };
+
+        for i in start..stop {
+            let item = self.getitem_by_idx(i as isize, vm)?;
+            if vm.bool_eq(&item, &value)? {
+                return Ok(i);
+            }
+        }
+        Err(vm.new_value_error("memoryview.index(x): x not in memoryview"))
+    }
+
     fn cast_to_1d(&self, format: PyStrRef, vm: &VirtualMachine) -> PyResult<Self> {
         let format_spec = Self::parse_format(format.as_str(), vm)?;
         let itemsize = format_spec.size();