Author: Armin Rigo <ar...@tunes.org>
Branch: py3.5
Changeset: r88804:cca0cbd9a9d6
Date: 2016-12-01 18:19 +0100
http://bitbucket.org/pypy/pypy/changeset/cca0cbd9a9d6/

Log:    small fixes in struct.iter_unpack()

diff --git a/pypy/module/struct/interp_struct.py 
b/pypy/module/struct/interp_struct.py
--- a/pypy/module/struct/interp_struct.py
+++ b/pypy/module/struct/interp_struct.py
@@ -105,7 +105,12 @@
 
 
 class W_UnpackIter(W_Root):
-    def __init__(self, w_struct, buf):
+    def __init__(self, space, w_struct, w_buffer):
+        buf = space.buffer_w(w_buffer, space.BUF_SIMPLE)
+        if buf.getlength() % w_struct.size != 0:
+            raise oefmt(get_error(space),
+                "iterative unpacking requires a bytes length multiple of %d",
+                w_struct.size)
         self.w_struct = w_struct
         self.buf = buf
         self.index = 0
@@ -159,8 +164,7 @@
         return unpack_from(space, jit.promote_string(self.format), w_buffer, 
offset)
 
     def descr_iter_unpack(self, space, w_buffer):
-        buf = space.buffer_w(w_buffer, space.BUF_SIMPLE)
-        return W_UnpackIter(self, buf)
+        return W_UnpackIter(space, self, w_buffer)
 
 W_Struct.typedef = TypeDef("Struct",
     __new__=interp2app(W_Struct.descr__new__.im_func),
@@ -174,15 +178,7 @@
     iter_unpack=interp2app(W_Struct.descr_iter_unpack),
 )
 
-@unwrap_spec(w_struct=W_Struct)
-def new_unpackiter(space, w_subtype, w_struct, w_buffer):
-    buf = space.buffer_w(w_buffer, space.BUF_SIMPLE)
-    w_res = space.allocate_instance(W_UnpackIter, w_subtype)
-    w_res.__init__(w_struct, buf)
-    return w_res
-
 W_UnpackIter.typedef = TypeDef("unpack_iterator",
-    __new__=interp2app(new_unpackiter),
     __iter__=interp2app(W_UnpackIter.descr_iter),
     __next__=interp2app(W_UnpackIter.descr_next),
     __length_hint__=interp2app(W_UnpackIter.descr_length_hint)
@@ -191,8 +187,7 @@
 @unwrap_spec(format=str)
 def iter_unpack(space, format, w_buffer):
     w_struct = W_Struct(space, format)
-    buf = space.buffer_w(w_buffer, space.BUF_SIMPLE)
-    return W_UnpackIter(w_struct, buf)
+    return W_UnpackIter(space, w_struct, w_buffer)
 
 def clearcache(space):
     """No-op on PyPy"""
diff --git a/pypy/module/struct/test/test_struct.py 
b/pypy/module/struct/test/test_struct.py
--- a/pypy/module/struct/test/test_struct.py
+++ b/pypy/module/struct/test/test_struct.py
@@ -403,6 +403,25 @@
         assert list(it) == [(0, 0), (0, 0)]
         it = self.struct.iter_unpack('ii', b)
         assert list(it) == [(0, 0), (0, 0)]
+        #
+        it = s.iter_unpack(b)
+        next(it)
+        assert it.__length_hint__() == 1
+        next(it)
+        assert it.__length_hint__() == 0
+        assert list(it) == []
+        assert it.__length_hint__() == 0
+
+    def test_iter_unpack_bad_length(self):
+        struct = self.struct
+        s = struct.Struct('!i')
+        lst = list(s.iter_unpack(b'1234'))
+        assert lst == [(0x31323334,)]
+        lst = list(s.iter_unpack(b''))
+        assert lst == []
+        raises(struct.error, s.iter_unpack, b'12345')
+        raises(struct.error, s.iter_unpack, b'123')
+        raises(struct.error, struct.iter_unpack, 'h', b'12345')
 
     def test___float__(self):
         class MyFloat(object):
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit

Reply via email to