Fix several struct module tests.

Author: skinnyBat

Lib/test/test_struct.py

@@ -320,8 +320,6 @@ def __int__(self):
             t = IntTester(format)
             t.run()
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_nN_code(self):
         # n and N don't exist in standard sizes
         def assertStructError(func, *args, **kwargs):
@@ -352,7 +350,7 @@ def test_p_code(self):
             self.assertEqual(got, expectedback)
 
     # TODO: RUSTPYTHON
-    @unittest.skip("")
+    @unittest.expectedFailure
     def test_705836(self):
         # SF bug 705836.  "<f" and ">f" had a severe rounding bug, where a carry
         # from the low-order discarded bits could propagate into the exponent
@@ -394,8 +392,6 @@ def test_1530559(self):
         self.assertRaises(struct.error, struct.pack, 'P', 1.0)
         self.assertRaises(struct.error, struct.pack, 'P', 1.5)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_unpack_from(self):
         test_string = b'abcd01234'
         fmt = '4s'
@@ -527,16 +523,15 @@ def __bool__(self):
         for c in [b'\x01', b'\x7f', b'\xff', b'\x0f', b'\xf0']:
             self.assertTrue(struct.unpack('>?', c)[0])
 
-    @unittest.skip("TODO: RUSTPYTHON")
+    # TODO: RUSTPYTHON
+    @unittest.expectedFailure
     def test_count_overflow(self):
         hugecount = '{}b'.format(sys.maxsize+1)
         self.assertRaises(struct.error, struct.calcsize, hugecount)
 
         hugecount2 = '{}b{}H'.format(sys.maxsize//2, sys.maxsize//2)
         self.assertRaises(struct.error, struct.calcsize, hugecount2)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_trailing_counter(self):
         store = array.array('b', b' '*100)
 
@@ -676,8 +671,6 @@ class UnpackIteratorTest(unittest.TestCase):
     Tests for iterative unpacking (struct.Struct.iter_unpack).
     """
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_construct(self):
         def _check_iterator(it):
             self.assertIsInstance(it, abc.Iterator)
@@ -720,8 +713,6 @@ def test_arbitrary_buffer(self):
         self.assertRaises(StopIteration, next, it)
         self.assertRaises(StopIteration, next, it)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_length_hint(self):
         lh = operator.length_hint
         s = struct.Struct('>IB')

vm/src/byteslike.rs

@@ -5,6 +5,7 @@ use crate::pyobject::{PyResult, TryFromObject, TypeProtocol};
 use crate::stdlib::array::{PyArray, PyArrayRef};
 use crate::vm::VirtualMachine;
 
+#[derive(Debug)]
 pub enum PyBytesLike {
     Bytes(PyBytesRef),
     Bytearray(PyByteArrayRef),
@@ -26,6 +27,18 @@ impl TryFromObject for PyBytesLike {
 }
 
 impl PyBytesLike {
+    pub fn len(&self) -> usize {
+        match self {
+            PyBytesLike::Bytes(b) => b.len(),
+            PyBytesLike::Bytearray(b) => b.borrow_value().len(),
+            PyBytesLike::Array(array) => array.len(),
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
     pub fn to_cow(&self) -> std::borrow::Cow<[u8]> {
         match self {
             PyBytesLike::Bytes(b) => b.get_value().into(),

vm/src/stdlib/pystruct.rs

@@ -21,7 +21,7 @@ mod _struct {
     use std::io::{Cursor, Read, Write};
     use std::iter::Peekable;
 
-    use crate::byteslike::PyBuffer;
+    use crate::byteslike::{PyBuffer, PyBytesLike};
     use crate::exceptions::PyBaseExceptionRef;
     use crate::function::Args;
     use crate::obj::{
@@ -100,11 +100,11 @@ mod _struct {
         fn parse(fmt: &str) -> Result<FormatSpec, String> {
             let mut chars = fmt.chars().peekable();
 
-            // First determine "<", ">","!" or "="
-            let endianness = parse_endiannes(&mut chars);
+            // First determine "@", "<", ">","!" or "="
+            let (is_native_size, endianness) = parse_size_and_endiannes(&mut chars);
 
             // Now, analyze struct string furter:
-            let codes = parse_format_codes(&mut chars)?;
+            let codes = parse_format_codes(&mut chars, is_native_size)?;
 
             Ok(FormatSpec { endianness, codes })
         }
@@ -190,36 +190,36 @@ mod _struct {
 
     /// Parse endianness
     /// See also: https://docs.python.org/3/library/struct.html?highlight=struct#byte-order-size-and-alignment
-    fn parse_endiannes<I>(chars: &mut Peekable<I>) -> Endianness
+    fn parse_size_and_endiannes<I>(chars: &mut Peekable<I>) -> (bool, Endianness)
     where
         I: Sized + Iterator<Item = char>,
     {
         match chars.peek() {
             Some('@') => {
                 chars.next().unwrap();
-                Endianness::Native
+                (true, Endianness::Native)
             }
             Some('=') => {
                 chars.next().unwrap();
-                Endianness::Native
+                (false, Endianness::Native)
             }
             Some('<') => {
                 chars.next().unwrap();
-                Endianness::Little
+                (false, Endianness::Little)
             }
             Some('>') => {
                 chars.next().unwrap();
-                Endianness::Big
+                (false, Endianness::Big)
             }
             Some('!') => {
                 chars.next().unwrap();
-                Endianness::Network
+                (false, Endianness::Network)
             }
-            _ => Endianness::Native,
+            _ => (true, Endianness::Native),
         }
     }
 
-    fn parse_format_codes<I>(chars: &mut Peekable<I>) -> Result<Vec<FormatCode>, String>
+    fn parse_format_codes<I>(chars: &mut Peekable<I>, is_native_size: bool) -> Result<Vec<FormatCode>, String>
     where
         I: Sized + Iterator<Item = char>,
     {
@@ -243,6 +243,9 @@ mod _struct {
             // determine format char:
             let c = chars.next();
             match c {
+                Some('n') | Some('N') if !is_native_size => {
+                    return Err("bad char in struct format".to_owned())
+                }
                 Some(c) if is_supported_format_character(c) => {
                     codes.push(FormatCode { repeat, code: c })
                 }
@@ -667,11 +670,11 @@ mod _struct {
     #[pyfunction]
     fn unpack(
         fmt: Either<PyStringRef, PyBytesRef>,
-        buffer: PyBytesRef,
+        buffer: PyBytesLike,
         vm: &VirtualMachine,
     ) -> PyResult<PyTuple> {
         let format_spec = FormatSpec::decode_and_parse(vm, &fmt)?;
-        format_spec.unpack(buffer.get_value(), vm)
+        buffer.with_ref(|buf| format_spec.unpack(buf, vm))
     }
 
     fn unpack_code<Endianness>(
@@ -721,26 +724,49 @@ mod _struct {
         Ok(())
     }
 
+    #[derive(FromArgs)]
+    struct UpdateFromArgs {
+        buffer: PyBytesLike,
+        #[pyarg(positional_or_keyword, default = "0")]
+        offset: isize
+    }
+
     #[pyfunction]
     fn unpack_from(
         fmt: Either<PyStringRef, PyBytesRef>,
-        buffer: PyBytesRef,
-        offset: isize,
+        args: UpdateFromArgs,
         vm: &VirtualMachine,
     ) -> PyResult<PyTuple> {
         let format_spec = FormatSpec::decode_and_parse(vm, &fmt)?;
-        let offset = get_buffer_offset(buffer.len(), offset, format_spec.size(), vm)?;
-        format_spec.unpack(&buffer.get_value()[offset..], vm)
+        let size = format_spec.size();
+        let offset = get_buffer_offset(args.buffer.len(), args.offset, size, vm)?;
+        args.buffer.with_ref(|buf| format_spec.unpack(&buf[offset..offset+size], vm))
     }
 
     #[pyclass(name = "unpack_iterator")]
     #[derive(Debug)]
     struct UnpackIterator {
         format_spec: FormatSpec,
-        buffer: PyBytesRef,
+        buffer: PyBytesLike,
         offset: AtomicCell<usize>,
     }
 
+    impl UnpackIterator {
+        fn new(vm: &VirtualMachine, format_spec: FormatSpec, buffer: PyBytesLike) -> PyResult<UnpackIterator> {
+            if format_spec.size() == 0 {
+                Err(new_struct_error(vm, "cannot iteratively unpack with a struct of length 0".to_owned()))
+            } else if buffer.len() % format_spec.size() != 0 {
+                Err(new_struct_error(vm, format!("iterative unpacking requires a buffer of a multiple of {} bytes", format_spec.size())))
+            } else {
+                Ok(UnpackIterator {
+                    format_spec,
+                    buffer,
+                    offset: AtomicCell::new(0),
+                })
+            }
+        }
+    }
+
     impl PyValue for UnpackIterator {
         fn class(vm: &VirtualMachine) -> PyClassRef {
             vm.class("_struct", "unpack_iterator")
@@ -756,24 +782,30 @@ mod _struct {
             if offset + size > self.buffer.len() {
                 Err(objiter::new_stop_iteration(vm))
             } else {
-                self.format_spec
-                    .unpack(&self.buffer.get_value()[offset..offset + size], vm)
+                self.buffer.with_ref(|buf| self.format_spec
+                    .unpack(&buf[offset..offset + size], vm))
             }
         }
+
+        #[pymethod(magic)]
+        fn iter(zelf: PyRef<Self>) -> PyRef<Self> {
+            zelf
+        }
+
+        #[pymethod(magic)]
+        fn length_hint(&self) -> usize {
+            self.buffer.len().saturating_sub(self.offset.load()) / self.format_spec.size()
+        }
     }
 
     #[pyfunction]
     fn iter_unpack(
         fmt: Either<PyStringRef, PyBytesRef>,
-        buffer: PyBytesRef,
+        buffer: PyBytesLike,
         vm: &VirtualMachine,
     ) -> PyResult<UnpackIterator> {
         let format_spec = FormatSpec::decode_and_parse(vm, &fmt)?;
-        Ok(UnpackIterator {
-            format_spec,
-            buffer,
-            offset: AtomicCell::new(0),
-        })
+        UnpackIterator::new(vm, format_spec, buffer)
     }
 
     #[pyfunction]
@@ -843,28 +875,24 @@ mod _struct {
         }
 
         #[pymethod]
-        fn unpack(&self, data: PyBytesRef, vm: &VirtualMachine) -> PyResult<PyTuple> {
-            self.spec.unpack(data.get_value(), vm)
+        fn unpack(&self, data: PyBytesLike, vm: &VirtualMachine) -> PyResult<PyTuple> {
+            data.with_ref(|buf| self.spec.unpack(buf, vm))
         }
 
         #[pymethod]
         fn unpack_from(
             &self,
-            buffer: PyBytesRef,
-            offset: isize,
+            args: UpdateFromArgs,
             vm: &VirtualMachine,
         ) -> PyResult<PyTuple> {
-            let offset = get_buffer_offset(buffer.len(), offset, self.size(), vm)?;
-            self.spec.unpack(&buffer.get_value()[offset..], vm)
+            let size = self.size();
+            let offset = get_buffer_offset(args.buffer.len(), args.offset, size, vm)?;
+            args.buffer.with_ref(|buf| self.spec.unpack(&buf[offset..offset+size], vm))
         }
 
         #[pymethod]
-        fn iter_unpack(&self, buffer: PyBytesRef) -> PyResult<UnpackIterator> {
-            Ok(UnpackIterator {
-                format_spec: self.spec.clone(),
-                buffer,
-                offset: AtomicCell::new(0),
-            })
+        fn iter_unpack(&self, buffer: PyBytesLike, vm: &VirtualMachine) -> PyResult<UnpackIterator> {
+            UnpackIterator::new(vm, self.spec.clone(), buffer)
         }
     }