RF+TST: add tests for en/decoding value in fields

Add tests for encoding / decoding numerical values in byte string
fields.

Refactor encoding.

Update comments / docstrings to note new encoding with numbers as ASCII
strings.

Author: matthew-brett

nibabel/streamlines/tests/test_trk.py

@@ -16,7 +16,7 @@
 from ..tractogram_file import HeaderError, HeaderWarning
 
 from .. import trk as trk_module
-from ..trk import TrkFile
+from ..trk import TrkFile, encode_value_in_name, decode_value_from_name
 from ..header import Field
 
 DATA = {}
@@ -468,3 +468,31 @@ def test_header_read_restore(self):
         assert_arr_dict_equal(TrkFile._read_header(bio), hdr_from_fname)
         # Check fileobject file position has not changed
         assert_equal(bio.tell(), hdr_pos)
+
+
+def test_encode_names():
+    # Test function for encoding numbers into property names
+    b0 = b'\x00'
+    assert_equal(encode_value_in_name(0, 'foo', 10),
+                 b'foo' + b0 * 7)
+    assert_equal(encode_value_in_name(1, 'foo', 10),
+                 b'foo' + b0 * 7)
+    assert_equal(encode_value_in_name(8, 'foo', 10),
+                 b'foo' + b0 + b'8' + b0 * 5)
+    assert_equal(encode_value_in_name(40, 'foobar', 10),
+                 b'foobar' + b0 + b'40' + b0)
+    assert_equal(encode_value_in_name(1, 'foobarbazz', 10), b'foobarbazz')
+    assert_raises(ValueError, encode_value_in_name, 1, 'foobarbazzz', 10)
+    assert_raises(ValueError, encode_value_in_name, 2, 'foobarbaz', 10)
+    assert_equal(encode_value_in_name(2, 'foobarba', 10), b'foobarba\x002')
+
+
+def test_decode_names():
+    # Test function for decoding name string into name, number
+    b0 = b'\x00'
+    assert_equal(decode_value_from_name(b''), ('', 0))
+    assert_equal(decode_value_from_name(b'foo' + b0 * 7), ('foo', 1))
+    assert_equal(decode_value_from_name(b'foo\x008' + b0 * 5), ('foo', 8))
+    assert_equal(decode_value_from_name(b'foobar\x0010\x00'), ('foobar', 10))
+    assert_raises(ValueError, decode_value_from_name, b'foobar\x0010\x01')
+    assert_raises(HeaderError, decode_value_from_name, b'foo\x0010\x00111')

nibabel/streamlines/trk.py

@@ -129,48 +129,46 @@ def get_affine_rasmm_to_trackvis(header):
 
 
 def encode_value_in_name(value, name, max_name_len=20):
-    """ Encodes a value in the last bytes of a string.
+    """ Return `name` as fixed-length string, appending `value` as string.
 
-    If `value` is one, then there is no encoding and the last bytes
-    are left untouched. Otherwise, a \x00 byte is added after `name`
-    and followed by the ascii represensation of the value.
+    Form output from `name` if `value <= 1` else `name` + ``\x00`` +
+    str(value).
 
-    This function also verifies that the length of name is less
-    than `max_name_len`.
+    Return output as fixed length string length `max_name_len`, padded with
+    ``\x00``.
+
+    This function also verifies that the modified length of name is less than
+    `max_name_len`.
 
     Parameters
     ----------
-    value : byte
-        Integer value between 0 and 255 to encode.
-    name : bytes
-        Name in which the last two bytes will serve to encode `value`.
+    value : int
+        Integer value to encode.
+    name : str
+        Name to which we may append an ascii / latin-1 representation of
+        `value`.
     max_name_len : int, optional
-        Maximum length name can have.
+        Maximum length of byte string that output can have.
 
     Returns
     -------
     encoded_name : bytes
-        Name containing the encoded value.
+        Name maybe followed by ``\x00`` and ascii / latin-1 representation of
+        `value`, padded with ``\x00`` bytes.
     """
-
     if len(name) > max_name_len:
         msg = ("Data information named '{0}' is too long"
                " (max {1} characters.)").format(name, max_name_len)
         raise ValueError(msg)
-    elif value > 1 and len(name) + len(str(value)) + 1 > max_name_len:
+    encoded_name = name if value <= 1 else name + '\x00' + str(value)
+    if len(encoded_name) > max_name_len:
         msg = ("Data information named '{0}' is too long (need to be less"
                " than {1} characters when storing more than one value"
                " for a given data information."
                ).format(name, max_name_len - (len(str(value)) + 1))
         raise ValueError(msg)
-
-    encoded_name = name
-    if value > 1:
-        # Store the name followed by \x00 and the `value` (in ascii).
-        encoded_name += '\x00' + str(value)
-
-    encoded_name = encoded_name.ljust(max_name_len, '\x00')
-    return encoded_name
+    # Fill to the end with zeros
+    return encoded_name.ljust(max_name_len, '\x00').encode('latin1')
 
 
 def decode_value_from_name(encoded_name):
@@ -388,7 +386,7 @@ def _read():
         return cls(tractogram, header=hdr)
 
     def save(self, fileobj):
-        """ Saves tractogram to a file-like object using TRK format.
+        """ Save tractogram to a filename or file-like object using TRK format.
 
         Parameters
         ----------
@@ -420,6 +418,7 @@ def save(self, fileobj):
             # Keep track of the beginning of the header.
             beginning = f.tell()
 
+            # Write temporary header that we will update at the end
             f.write(header.tostring())
 
             i4_dtype = np.dtype("<i4")  # Always save in little-endian.
@@ -449,8 +448,8 @@ def save(self, fileobj):
             property_name = np.zeros(MAX_NB_NAMED_PROPERTIES_PER_STREAMLINE,
                                      dtype='S20')
             for i, name in enumerate(data_for_streamline_keys):
-                # Use the last two bytes of the name to store the number of
-                # values associated to this data_for_streamline.
+                # Append number of values as ascii to zero-terminated name
+                # to encode number of values into trackvis name.
                 nb_values = data_for_streamline[name].shape[-1]
                 property_name[i] = encode_value_in_name(nb_values, name)
             header['property_name'][:] = property_name
@@ -466,8 +465,8 @@ def save(self, fileobj):
             data_for_points_keys = sorted(data_for_points.keys())
             scalar_name = np.zeros(MAX_NB_NAMED_SCALARS_PER_POINT, dtype='S20')
             for i, name in enumerate(data_for_points_keys):
-                # Use the last two bytes of the name to store the number of
-                # values associated to this data_for_streamline.
+                # Append number of values as ascii to zero-terminated name
+                # to encode number of values into trackvis name.
                 nb_values = data_for_points[name].shape[-1]
                 scalar_name[i] = encode_value_in_name(nb_values, name)
             header['scalar_name'][:] = scalar_name