Commits

Brian Kearns committed 99632e3

properly use buffers in struct.unpack_from

  • Participants
  • Parent commits ab7d9bd

Comments (0)

Files changed (2)

File pypy/module/struct/formatiterator.py

 
 
 class PackFormatIterator(FormatIterator):
-
     def __init__(self, space, args_w, size):
         self.space = space
         self.args_w = args_w
 
 
 class UnpackFormatIterator(FormatIterator):
-
-    def __init__(self, space, input):
+    def __init__(self, space, buf):
         self.space = space
-        self.input = input
-        self.inputpos = 0
+        self.buf = buf
+        self.length = buf.getlength()
+        self.pos = 0
         self.result_w = []     # list of wrapped objects
 
     # See above comment on operate.
     _operate_is_specialized_ = True
 
     def align(self, mask):
-        self.inputpos = (self.inputpos + mask) & ~mask
+        self.pos = (self.pos + mask) & ~mask
 
     def finished(self):
-        if self.inputpos != len(self.input):
+        if self.pos != self.length:
             raise StructError("unpack str size too long for format")
 
     def read(self, count):
-        end = self.inputpos + count
-        if end > len(self.input):
+        end = self.pos + count
+        if end > self.length:
             raise StructError("unpack str size too short for format")
-        s = self.input[self.inputpos : end]
-        self.inputpos = end
+        s = self.buf.getslice(self.pos, end, 1, end - self.pos)
+        self.pos = end
         return s
 
     @specialize.argtype(1)

File pypy/module/struct/interp_struct.py

 from rpython.rlib import jit
+from rpython.rlib.buffer import SubBuffer
 from rpython.rlib.rstruct.error import StructError, StructOverflowError
 from rpython.rlib.rstruct.formatiterator import CalcSizeFormatIterator
 from rpython.tool.sourcetools import func_with_new_name
     buf.setslice(offset, res)
 
 
-@unwrap_spec(format=str, input='bufferstr')
-def unpack(space, format, input):
-    fmtiter = UnpackFormatIterator(space, input)
+def _unpack(space, format, buf):
+    fmtiter = UnpackFormatIterator(space, buf)
     try:
         fmtiter.interpret(format)
     except StructOverflowError, e:
     return space.newtuple(fmtiter.result_w[:])
 
 
-# XXX inefficient
+@unwrap_spec(format=str)
+def unpack(space, format, w_str):
+    buf = space.getarg_w('s*', w_str)
+    return _unpack(space, format, buf)
+
+
 @unwrap_spec(format=str, offset=int)
-def unpack_from(space, format, w_buf, offset=0):
+def unpack_from(space, format, w_buffer, offset=0):
     size = _calcsize(space, format)
-    buf = space.getarg_w('z*', w_buf)
+    buf = space.getarg_w('z*', w_buffer)
     if buf is None:
         w_module = space.getbuiltinmodule('struct')
         w_error = space.getattr(w_module, space.wrap('error'))
         raise oefmt(w_error,
                     "unpack_from requires a buffer of at least %d bytes",
                     size)
-    data = buf.getslice(offset, offset + size, 1, size)
-    return unpack(space, format, data)
+    buf = SubBuffer(buf, offset, size)
+    return _unpack(space, format, buf)
 
 
 class W_Struct(W_Root):