RF - moved extensions into header; fixed __eq__ bug with tests

Author: matthew-brett

nibabel/nifti1.py

@@ -366,6 +366,8 @@ def __repr__(self):
         return s
 
     def __eq__(self, other):
+        if len(self) != len(other):
+            return False
         for i, e in enumerate(self):
             if not e == other[i]:
                 return False
@@ -389,8 +391,9 @@ def write_to(self, fileobj, byteswap):
         '''
         # not extensions -> nothing to do
         if not len(self):
+            # no extensions: be nice and write appropriate flag
+            fileobj.write(np.array((0, 0, 0, 0), dtype=np.int8).tostring())
             return
-
         # since we have extensions write the appropriate flag
         fileobj.write(np.array((1, 0, 0, 0), dtype=np.int8).tostring())
         # and now each extension
@@ -428,12 +431,10 @@ def from_fileobj(klass, fileobj, size, byteswap):
             extension_status = np.fromstring(extension_status, dtype=np.int8)
             if byteswap:
                 extension_status = extension_status.byteswap()
-
         # NIfTI1 says: if first element is non-zero there are extensions present
         # if not there is nothing left to do
         if not extension_status[0]:
             return extensions
-
         # note that we read the extension flag
         if not size < 0:
             size = size - 4
@@ -481,7 +482,11 @@ def from_fileobj(klass, fileobj, size, byteswap):
 
 
 class Nifti1Header(SpmAnalyzeHeader):
-    ''' Class for NIFTI1 header '''
+    ''' Class for NIFTI1 header 
+    
+    The NIFTI1 header has many more coded fields than the simpler Analyze
+    variants.  Analyze headers also have extensions
+    '''
     # Copies of module level definitions
     _dtype = header_dtype
     _data_type_codes = data_type_codes
@@ -497,6 +502,60 @@ class Nifti1Header(SpmAnalyzeHeader):
     has_data_slope = True
     has_data_intercept = True
 
+    # Extension class; should implement __call__ for contruction, and
+    # ``from_fileobj`` for reading from file
+    exts_klass = Nifti1Extensions
+
+    def __init__(self,
+                 binaryblock=None,
+                 endianness=None,
+                 check=True,
+                 extensions=()):
+        ''' Initialize header from binary data block and extensions
+        '''
+        super(Nifti1Header, self).__init__(binaryblock,
+                                           endianness,
+                                           check)
+        self.extensions = self.exts_klass(extensions)
+
+    def copy(self):
+        ''' Return copy of header
+
+        Take extensions as well as header
+        '''
+        return self.__class__(
+            self.binaryblock,
+            self.endianness, 
+            False,
+            self.extensions)
+
+    @classmethod
+    def from_fileobj(klass, fileobj, endianness=None, check=True):
+        raw_str = fileobj.read(klass._dtype.itemsize)
+        hdr = klass(raw_str, endianness, check)
+        hdr_len = hdr._header_len()
+        if hdr_len == -1:
+            extsize = -1
+        else:
+            extsize = hdr_len - fileobj.tell()
+        byteswap = endian_codes['native'] != hdr.endianness
+        hdr.extensions = klass.exts_klass.from_fileobj(fileobj, extsize, byteswap)
+        return hdr
+
+    def write_to(self, fileobj):
+        super(Nifti1Header, self).write_to(fileobj)
+        byteswap = endian_codes['native'] != self.endianness
+        self.extensions.write_to(fileobj, byteswap)
+
+    def _header_len(self):
+        ''' Return header length in bytes or -1 for unknown
+
+        This will be -1 for headers that are their own files, as in the .hdr
+        file of a nifti pair, or the same as the start of the data (vox_offset)
+        in single file niftis
+        '''
+        return self._header_data['vox_offset']
+
     def get_best_affine(self):
         ''' Select best of available transforms '''
         hdr = self._header_data
@@ -1278,6 +1337,14 @@ def _set_format_specifics(self):
         ''' Utility routine to set format specific header stuff '''
         self._header_data['magic'] = 'ni1'
 
+    def _header_len(self):
+        ''' Return header length in bytes or -1 for unknown
+
+        This will be -1 for headers that are their own files, as in the .hdr
+        file of a nifti pair, or the same as the start of the data (vox_offset)
+        in single file niftis
+        '''
+        return -1
 
 class Nifti1Pair(analyze.AnalyzeImage):
     header_class = Nifti1PairHeader
@@ -1287,20 +1354,6 @@ def from_file_map(klass, file_map):
         hdrf, imgf = klass._get_open_files(file_map, 'rb')
         header = klass.header_class.from_fileobj(hdrf)
         extra = None
-        # handle extensions
-        # assume the fileptr is just after header (magic field)
-        # determine how much to read when parsing the extensions
-        if header['vox_offset'] == 0:
-            # read till the end of the header
-            extsize = -1
-        else:
-            extsize = header['vox_offset'] - hdrf.tell()
-        extensions = Nifti1Extensions.from_fileobj(
-                        hdrf, extsize,
-                        endian_codes['native'] != header.endianness)
-        # XXX maybe always do that?
-        if len(extensions):
-            extra = {'extensions': extensions}
         affine = header.get_best_affine()
         hdr_copy = header.copy()
         data = klass.ImageArrayProxy(imgf, hdr_copy)
@@ -1315,13 +1368,6 @@ def _write_header(self, header_file, header, slope, inter):
                                               header,
                                               slope,
                                               inter)
-        if not self.extra.has_key('extensions'):
-            # no extensions: be nice and write appropriate flag
-            header_file.write(np.array((0, 0, 0, 0), dtype=np.int8).tostring())
-        else:
-            self.extra['extensions'].write_to(
-                        header_file,
-                        endian_codes['native'] != header.endianness)
 
     def update_header(self):
         ''' Harmonize header with image data and affine

nibabel/tests/test_nifti1.py

@@ -18,8 +18,8 @@
 from ..spatialimages import HeaderDataError
 from .. import nifti1 as nifti1
 from ..nifti1 import (load, Nifti1Header, Nifti1Image,
-                      Nifti1Pair, Nifti1Extension, data_type_codes,
-                      extension_codes, slice_order_codes)
+                      Nifti1Pair, Nifti1Extension, Nifti1Extensions,
+                      data_type_codes, extension_codes, slice_order_codes)
 
 from numpy.testing import assert_array_equal, assert_array_almost_equal
 from nose.tools import (assert_true, assert_false, assert_equal,
@@ -387,73 +387,85 @@ def test_extension_codes():
         ext = Nifti1Extension(k, 'somevalue')
 
 
-@parametric
+def test_extension_list():
+    ext_c0 = Nifti1Extensions()
+    ext_c1 = Nifti1Extensions()
+    assert_equal(ext_c0, ext_c1)
+    ext = Nifti1Extension('comment', '123')
+    ext_c1.append(ext)
+    assert_false(ext_c0 == ext_c1)
+    ext_c0.append(ext)
+    assert_true(ext_c0 == ext_c1)
+
+
 def test_nifti_extensions():
     nim = load(image_file)
     # basic checks of the available extensions
-    ext = nim.extra['extensions']
-    yield assert_true(len(ext) == 2)
-    yield assert_true(ext.count('comment') == 2)
-    yield assert_true(ext.count('afni') == 0)
-    yield assert_true(ext.get_codes() == [6, 6])
-    yield assert_true((ext.get_sizeondisk() - 4) % 16 == 0)
+    hdr = nim.get_header()
+    exts_container = hdr.extensions
+    assert_true(len(exts_container) == 2)
+    assert_true(exts_container.count('comment') == 2)
+    assert_true(exts_container.count('afni') == 0)
+    assert_true(exts_container.get_codes() == [6, 6])
+    assert_true((exts_container.get_sizeondisk() - 4) % 16 == 0)
     # first extension should be short one
-    yield assert_true(ext[0].get_content() == 'extcomment1')
+    assert_true(exts_container[0].get_content() == 'extcomment1')
     # add one
     afniext = Nifti1Extension('afni', '<xml></xml>')
-    ext.append(afniext)
-    yield assert_true(ext.get_codes() == [6, 6, 4])
-    yield assert_true(ext.count('comment') == 2)
-    yield assert_true(ext.count('afni') == 1)
-    yield assert_true((ext.get_sizeondisk() - 4) % 16 == 0)
+    exts_container.append(afniext)
+    assert_true(exts_container.get_codes() == [6, 6, 4])
+    assert_true(exts_container.count('comment') == 2)
+    assert_true(exts_container.count('afni') == 1)
+    assert_true((exts_container.get_sizeondisk() - 4) % 16 == 0)
     # delete one
-    del ext[1]
-    yield assert_true(ext.get_codes() == [6, 4])
-    yield assert_true(ext.count('comment') == 1)
-    yield assert_true(ext.count('afni') == 1)
+    del exts_container[1]
+    assert_true(exts_container.get_codes() == [6, 4])
+    assert_true(exts_container.count('comment') == 1)
+    assert_true(exts_container.count('afni') == 1)
 
 
-@parametric
 def test_loadsave_cycle():
     nim = load(image_file)
     # ensure we have extensions
-    yield assert_true(nim.extra.has_key('extensions'))
-    yield assert_true(len(nim.extra['extensions']))
+    hdr = nim.get_header()
+    exts_container = hdr.extensions
+    assert_true(len(exts_container) > 0)
     # write into the air ;-)
     stio = StringIO()
     nim.file_map['image'].fileobj = stio
     nim.to_file_map()
     stio.seek(0)
     # reload
     lnim = Nifti1Image.from_file_map(nim.file_map)
-    yield assert_true(lnim.extra.has_key('extensions'))
-    yield assert_equal(nim.extra['extensions'],
-                       lnim.extra['extensions'])
+    hdr = lnim.get_header()
+    lexts_container = hdr.extensions
+    assert_equal(exts_container,
+                 lexts_container)
     # build int16 image
     data = np.ones((2,3,4,5), dtype='int16')
     img = Nifti1Image(data, np.eye(4))
     hdr = img.get_header()
-    yield assert_equal(hdr.get_data_dtype(), np.int16)
+    assert_equal(hdr.get_data_dtype(), np.int16)
     # default should have no scaling
-    yield assert_equal(hdr.get_slope_inter(), (1.0, 0.0))
+    assert_equal(hdr.get_slope_inter(), (1.0, 0.0))
     # set scaling
     hdr.set_slope_inter(2, 8)
-    yield assert_equal(hdr.get_slope_inter(), (2, 8))
+    assert_equal(hdr.get_slope_inter(), (2, 8))
     # now build new image with updated header
     wnim = Nifti1Image(data, np.eye(4), header=hdr)
-    yield assert_equal(wnim.get_data_dtype(), np.int16)
-    yield assert_equal(wnim.get_header().get_slope_inter(), (2, 8))
+    assert_equal(wnim.get_data_dtype(), np.int16)
+    assert_equal(wnim.get_header().get_slope_inter(), (2, 8))
     # write into the air again ;-)
     stio = StringIO()
     wnim.file_map['image'].fileobj = stio
     wnim.to_file_map()
     stio.seek(0)
     lnim = Nifti1Image.from_file_map(wnim.file_map)
-    yield assert_equal(lnim.get_data_dtype(), np.int16)
+    assert_equal(lnim.get_data_dtype(), np.int16)
     # the test below does not pass, because the slope and inter are
     # always reset from the data, by the image write
     raise SkipTest
-    yield assert_equal(lnim.get_header().get_slope_inter(), (2, 8))
+    assert_equal(lnim.get_header().get_slope_inter(), (2, 8))
 
 
 @parametric