img: add _MarkerFinder.next()

Steve Canny · Steve Canny · commit 17610b091844 · 2014-02-27T12:08:35.000-08:00
diff --git a/docx/image/helpers.py b/docx/image/helpers.py
@@ -25,6 +25,12 @@ def __init__(self, stream, byte_order, base_offset=0):
         )
         self._base_offset = base_offset
 
+    def read(self, count):
+        """
+        Allow pass-through read() call
+        """
+        return self._stream.read(count)
+
     def read_byte(self, base=None, offset=0):
         """
         Return the int value of the byte at the file position defined by
@@ -61,6 +67,12 @@ def seek(self, base, offset=0):
         location = self._base_offset + base + offset
         self._stream.seek(location)
 
+    def tell(self):
+        """
+        Allow pass-through tell() call
+        """
+        return self._stream.tell()
+
     def _read_bytes(self, byte_count, base, offset):
         self.seek(base, offset)
         bytes_ = self._stream.read(byte_count)
diff --git a/docx/image/jpeg.py b/docx/image/jpeg.py
@@ -158,7 +158,56 @@ def next(self, start):
         following the 2-byte marker code, the start of the marker segment,
         for those markers that have a segment.
         """
-        raise NotImplementedError
+        position = start
+        while True:
+            # skip over any non-\xFF bytes
+            position = self._offset_of_next_ff_byte(start=position)
+            # skip over any \xFF padding bytes
+            position, byte_ = self._next_non_ff_byte(start=position+1)
+            # 'FF 00' sequence is not a marker, start over if found
+            if byte_ == b'\x00':
+                continue
+            # this is a marker, gather return values and break out of scan
+            marker_code, segment_offset = byte_, position+1
+            break
+        return marker_code, segment_offset
+
+    def _next_non_ff_byte(self, start):
+        """
+        Return an offset, byte 2-tuple for the next byte in *stream* that is
+        not '\xFF', starting with the byte at offset *start*. If the byte at
+        offset *start* is not '\xFF', *start* and the returned *offset* will
+        be the same.
+        """
+        self._stream.seek(start)
+        byte_ = self._read_byte()
+        while byte_ == b'\xFF':
+            byte_ = self._read_byte()
+        offset_of_non_ff_byte = self._stream.tell() - 1
+        return offset_of_non_ff_byte, byte_
+
+    def _offset_of_next_ff_byte(self, start):
+        """
+        Return the offset of the next '\xFF' byte in *stream* starting with
+        the byte at offset *start*. Returns *start* if the byte at that
+        offset is a hex 255; it does not necessarily advance in the stream.
+        """
+        self._stream.seek(start)
+        byte_ = self._read_byte()
+        while byte_ != b'\xFF':
+            byte_ = self._read_byte()
+        offset_of_ff_byte = self._stream.tell() - 1
+        return offset_of_ff_byte
+
+    def _read_byte(self):
+        """
+        Return the next byte read from stream. Raise Exception if stream is
+        at end of file.
+        """
+        byte_ = self._stream.read(1)
+        if not byte_:
+            raise Exception('unexpected end of file')
+        return byte_
 
 
 def _MarkerFactory(marker_code, stream, offset):
diff --git a/tests/image/test_jpeg.py b/tests/image/test_jpeg.py
@@ -159,6 +159,11 @@ def it_can_construct_from_a_stream(self, from_stream_fixture):
         _MarkerFinder__init_.assert_called_once_with(stream_)
         assert isinstance(marker_finder, _MarkerFinder)
 
+    def it_can_find_the_next_marker_after_a_given_offset(self, next_fixture):
+        marker_finder, start, expected_code_and_offset = next_fixture
+        marker_code, segment_offset = marker_finder.next(start)
+        assert (marker_code, segment_offset) == expected_code_and_offset
+
     # fixtures -------------------------------------------------------
 
     @pytest.fixture
@@ -169,6 +174,23 @@ def from_stream_fixture(self, stream_, _MarkerFinder__init_):
     def _MarkerFinder__init_(self, request):
         return initializer_mock(request, _MarkerFinder)
 
+    @pytest.fixture(params=[
+        (0, JPEG_MARKER_CODE.SOI,   2),
+        (1, JPEG_MARKER_CODE.APP0,  4),
+        (2, JPEG_MARKER_CODE.APP0,  4),
+        (3, JPEG_MARKER_CODE.EOI,  12),
+        (4, JPEG_MARKER_CODE.EOI,  12),
+        (6, JPEG_MARKER_CODE.EOI,  12),
+        (8, JPEG_MARKER_CODE.EOI,  12),
+    ])
+    def next_fixture(self, request):
+        start, marker_code, segment_offset = request.param
+        bytes_ = b'\xFF\xD8\xFF\xE0\x00\x01\xFF\x00\xFF\xFF\xFF\xD9'
+        stream_reader = StreamReader(BytesIO(bytes_), BIG_ENDIAN)
+        marker_finder = _MarkerFinder(stream_reader)
+        expected_code_and_offset = (marker_code, segment_offset)
+        return marker_finder, start, expected_code_and_offset
+
     @pytest.fixture
     def stream_(self, request):
         return instance_mock(request, BytesIO)
