Merge pull request #334 from matthew-brett/zero-terminal-strings

MRG: add routine to read zero-terminal strings

Routine to read zero-terminal strings of arbitrary length from binary file
stream. Inspired by Thomas Emmerling's readCstring function in the 
BrainVoyager PR.

Author: matthew-brett

nibabel/fileutils.py

@@ -0,0 +1,62 @@
+# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+# ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+#
+#   See COPYING file distributed along with the NiBabel package for the
+#   copyright and license terms.
+#
+# ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+""" Utilities for reading and writing to binary file formats
+"""
+
+def read_zt_byte_strings(fobj, n_strings=1, bufsize=1024):
+    """Read zero-terminated byte strings from a file object `fobj`
+
+    Returns byte strings with terminal zero stripped.
+
+    Found strings can be of any length.
+
+    The file position of `fobj` on exit will be at the byte after the terminal
+    0 of the final read byte string.
+
+    Parameters
+    ----------
+    f : fileobj
+        File object to use.  Should implement ``read``, returning byte objects
+        (str in Python 2), and ``seek(n, 1)`` to seek from current file
+        position.
+    n_strings : int, optional
+        Number of byte strings to return
+    bufsize: int, optional
+       Define chunk size to load from file while searching for zero terminals.
+       We load this many bytes at a time from the file, but the returned
+       strings can be longer than `bufsize`.
+
+    Returns
+    -------
+    byte_strings : list
+        List of byte strings, where strings do not include the terminal 0
+    """
+    byte_strings = []
+    trailing = b''
+    while True:
+        buf = fobj.read(bufsize)
+        eof = len(buf) < bufsize  # end of file
+        zt_strings = buf.split(b'\x00')
+        if len(zt_strings) > 1:  # At least one 0
+            byte_strings += [trailing + zt_strings[0]] + zt_strings[1:-1]
+            trailing = zt_strings[-1]
+        else:  # No 0
+            trailing += zt_strings[0]
+        n_found = len(byte_strings)
+        if eof or n_found >= n_strings:
+            break
+    if n_found < n_strings:
+        raise ValueError('Expected {0} strings, found {1}'.format(
+            n_strings, n_found))
+    n_extra = n_found - n_strings
+    leftover_strings = byte_strings[n_strings:] + [trailing]
+    # Add number of extra strings to account for lost terminal 0s
+    extra_bytes = sum(len(bs) for bs in leftover_strings) + n_extra
+    fobj.seek(-extra_bytes, 1)  # seek back from current position
+    return byte_strings[:n_strings]

nibabel/tests/test_fileutils.py

@@ -0,0 +1,56 @@
+# emacs: -*- mode: python-mode; py-indent-offset: 4; indent-tabs-mode: nil -*-
+# vi: set ft=python sts=4 ts=4 sw=4 et:
+# ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+#
+#   See COPYING file distributed along with the NiBabel package for the
+#   copyright and license terms.
+#
+# ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
+""" Testing fileutils module
+"""
+
+import numpy as np
+
+from ..fileutils import read_zt_byte_strings
+
+from numpy.testing import (assert_almost_equal,
+                           assert_array_equal)
+
+from nose.tools import (assert_true, assert_false, assert_raises,
+                        assert_equal, assert_not_equal)
+
+
+from ..tmpdirs import InTemporaryDirectory
+
+
+def test_read_zt_byte_strings():
+    # sample binary block
+    binary = b'test.fmr\x00test.prt\x00something'
+    with InTemporaryDirectory():
+        # create a tempfile
+        path = 'test.bin'
+        fwrite = open(path, 'wb')
+        # write the binary block to it
+        fwrite.write(binary)
+        fwrite.close()
+        # open it again
+        fread = open(path, 'rb')
+        # test readout of one string
+        assert_equal(read_zt_byte_strings(fread), [b'test.fmr'])
+        # test new file position
+        assert_equal(fread.tell(), 9)
+        # manually rewind
+        fread.seek(0)
+        # test readout of two strings
+        assert_equal(read_zt_byte_strings(fread, 2),
+                     [b'test.fmr', b'test.prt'])
+        assert_equal(fread.tell(), 18)
+        # test readout of more strings than present
+        fread.seek(0)
+        assert_raises(ValueError, read_zt_byte_strings, fread, 3)
+        fread.seek(9)
+        assert_raises(ValueError, read_zt_byte_strings, fread, 2)
+        # Try with a small bufsize
+        fread.seek(0)
+        assert_equal(read_zt_byte_strings(fread, 2, 4),
+                     [b'test.fmr', b'test.prt'])