img: add StreamReader.read_long()

Steve Canny · Steve Canny · commit 2c6d96aa6141 · 2014-02-27T12:08:34.000-08:00
diff --git a/docx/image/helpers.py b/docx/image/helpers.py
@@ -25,6 +25,16 @@ def __init__(self, stream, byte_order, base_offset=0):
         )
         self._base_offset = base_offset
 
+    def read_long(self, base=None, offset=0):
+        """
+        Return the int value of the four bytes at the file position defined by
+        self._base_offset + *base* + *offset*. If *base* is None, the long is
+        read from the current position in the stream. The endian setting of
+        this instance is used to interpret the byte layout of the long.
+        """
+        fmt = '<L' if self._byte_order is LITTLE_ENDIAN else '>L'
+        return self._read_int(fmt, base, offset)
+
     def read_str(self, char_count, base, offset=0):
         """
         Return a string containing the *char_count* bytes at the file
@@ -49,6 +59,12 @@ def _read_bytes(self, byte_count, base, offset):
             raise UnexpectedEndOfFileError
         return bytes_
 
+    def _read_int(self, fmt, base, offset):
+        struct = Struct(fmt)
+        if base is None:
+            base = self._stream.tell() - self._base_offset
+        return self._unpack_item(struct, base, offset)
+
     def _unpack_item(self, struct, base, offset):
         bytes_ = self._read_bytes(struct.size, base, offset)
         return struct.unpack(bytes_)[0]
diff --git a/tests/image/test_helpers.py b/tests/image/test_helpers.py
@@ -10,7 +10,7 @@
 
 from docx.compat import BytesIO
 from docx.image.exceptions import UnexpectedEndOfFileError
-from docx.image.helpers import BIG_ENDIAN, StreamReader
+from docx.image.helpers import BIG_ENDIAN, LITTLE_ENDIAN, StreamReader
 
 
 class DescribeStreamReader(object):
@@ -26,8 +26,23 @@ def it_raises_on_unexpected_EOF(self, read_str_fixture):
         with pytest.raises(UnexpectedEndOfFileError):
             stream_rdr.read_str(9, 2)
 
+    def it_can_read_a_long(self, read_long_fixture):
+        stream_rdr, offset, expected_int = read_long_fixture
+        l = stream_rdr.read_long(offset)
+        assert l == expected_int
+
     # fixtures -------------------------------------------------------
 
+    @pytest.fixture(params=[
+        (BIG_ENDIAN,    b'\xBE\x00\x00\x00\x2A\xEF', 1, 42),
+        (LITTLE_ENDIAN, b'\xBE\xEF\x2A\x00\x00\x00', 2, 42),
+    ])
+    def read_long_fixture(self, request):
+        byte_order, bytes_, offset, expected_int = request.param
+        stream = BytesIO(bytes_)
+        stream_rdr = StreamReader(stream, byte_order)
+        return stream_rdr, offset, expected_int
+
     @pytest.fixture
     def read_str_fixture(self):
         stream = BytesIO(b'\x01\x02foobar\x03\x04')
