img: add StreamReader.read_str()

Steve Canny · Steve Canny · commit b5da364a8eec · 2014-02-27T12:08:34.000-08:00
diff --git a/docx/image/exceptions.py b/docx/image/exceptions.py
@@ -7,6 +7,12 @@ class InvalidImageStreamError(Exception):
     """
 
 
+class UnexpectedEndOfFileError(Exception):
+    """
+    EOF was unexpectedly encountered while reading an image stream.
+    """
+
+
 class UnrecognizedImageError(Exception):
     """
     The provided image stream could not be recognized.
diff --git a/docx/image/helpers.py b/docx/image/helpers.py
@@ -2,9 +2,13 @@
 
 from __future__ import absolute_import, division, print_function
 
+from struct import Struct
 
-_BIG_ENDIAN = '>'
-_LITTLE_ENDIAN = '<'
+from .exceptions import UnexpectedEndOfFileError
+
+
+BIG_ENDIAN = '>'
+LITTLE_ENDIAN = '<'
 
 
 class StreamReader(object):
@@ -17,6 +21,34 @@ def __init__(self, stream, byte_order, base_offset=0):
         super(StreamReader, self).__init__()
         self._stream = stream
         self._byte_order = (
-            _LITTLE_ENDIAN if byte_order == _LITTLE_ENDIAN else _BIG_ENDIAN
+            LITTLE_ENDIAN if byte_order == LITTLE_ENDIAN else BIG_ENDIAN
         )
         self._base_offset = base_offset
+
+    def read_str(self, char_count, base, offset=0):
+        """
+        Return a string containing the *char_count* bytes at the file
+        position determined by self._base_offset + *base* + *offset*.
+        """
+        def str_struct(char_count):
+            format_ = '%ds' % char_count
+            return Struct(format_)
+        struct = str_struct(char_count)
+        chars = self._unpack_item(struct, base, offset)
+        unicode_str = chars.decode('UTF-8')
+        return unicode_str
+
+    def seek(self, base, offset=0):
+        location = self._base_offset + base + offset
+        self._stream.seek(location)
+
+    def _read_bytes(self, byte_count, base, offset):
+        self.seek(base, offset)
+        bytes_ = self._stream.read(byte_count)
+        if len(bytes_) < byte_count:
+            raise UnexpectedEndOfFileError
+        return bytes_
+
+    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
@@ -6,6 +6,31 @@
 
 from __future__ import absolute_import, print_function
 
+import pytest
+
+from docx.compat import BytesIO
+from docx.image.exceptions import UnexpectedEndOfFileError
+from docx.image.helpers import BIG_ENDIAN, StreamReader
+
 
 class DescribeStreamReader(object):
-    pass
+
+    def it_can_read_a_string_of_specified_len_at_offset(
+            self, read_str_fixture):
+        stream_rdr, expected_string = read_str_fixture
+        s = stream_rdr.read_str(6, 2)
+        assert s == 'foobar'
+
+    def it_raises_on_unexpected_EOF(self, read_str_fixture):
+        stream_rdr = read_str_fixture[0]
+        with pytest.raises(UnexpectedEndOfFileError):
+            stream_rdr.read_str(9, 2)
+
+    # fixtures -------------------------------------------------------
+
+    @pytest.fixture
+    def read_str_fixture(self):
+        stream = BytesIO(b'\x01\x02foobar\x03\x04')
+        stream_rdr = StreamReader(stream, BIG_ENDIAN)
+        expected_string = 'foobar'
+        return stream_rdr, expected_string
