Modify save, remove vestigates of class_map / ext_map

Author: Ben Cipollini

nibabel/__init__.py

@@ -61,7 +61,7 @@
 from .orientations import (io_orientation, orientation_affine,
                            flip_axis, OrientationError,
                            apply_orientation, aff2axcodes)
-from .imageclasses import class_map, ext_map
+from .imageclasses import class_map, ext_map, all_image_classes
 from . import trackvis
 from . import mriutils
 

nibabel/loadsave.py

@@ -13,10 +13,8 @@
 
 from .filename_parser import splitext_addext
 from .volumeutils import BinOpener
-from .nifti1 import Nifti1Image, Nifti1Pair
-from .nifti2 import Nifti2Image, Nifti2Pair
 from .spatialimages import ImageFileError
-from .imageclasses import class_map, ext_map, all_image_classes
+from .imageclasses import all_image_classes
 from .arrayproxy import is_proxy
 
 
@@ -60,14 +58,22 @@ def save(img, filename):
     -------
     None
     '''
+
+    # Save the type as expected
     try:
         img.to_filename(filename)
     except ImageFileError:
         pass
     else:
         return
-    froot, ext, trailing = splitext_addext(filename, ('.gz', '.bz2'))
+
+    # Be nice to users by making common implicit conversions
+    froot, ext, trailing = splitext_addext(filename, img._compressed_exts)
+    lext = ext.lower()
+
     # Special-case Nifti singles and Pairs
+    from .nifti1 import Nifti1Image, Nifti1Pair  # Inline imports, as this file
+    from .nifti2 import Nifti2Image, Nifti2Pair  # really shouldn't reference any image type
     if type(img) == Nifti1Image and ext in ('.img', '.hdr'):
         klass = Nifti1Pair
     elif type(img) == Nifti2Image and ext in ('.img', '.hdr'):
@@ -76,9 +82,14 @@ def save(img, filename):
         klass = Nifti1Image
     elif type(img) == Nifti2Pair and ext == '.nii':
         klass = Nifti2Image
-    else:
-        img_type = ext_map[ext]
-        klass = class_map[img_type]['class']
+    else:  # arbitrary conversion
+        valid_klasses = filter(lambda klass: klass.is_valid_extension(lext),
+                               all_image_classes)
+        if len(valid_klasses) > 0:
+            klass = valid_klasses[0]
+        else:
+            raise ImageFileError('Cannot work out file type of "%s"' %
+                                 filename)
     converted = klass.from_image(img)
     converted.to_filename(filename)
 

nibabel/tests/test_files_interface.py

@@ -12,7 +12,7 @@
 
 import numpy as np
 
-from .. import class_map, Nifti1Image, Nifti1Pair, MGHImage
+from .. import Nifti1Image, Nifti1Pair, MGHImage, all_image_classes
 from ..externals.six import BytesIO
 from ..fileholders import FileHolderError
 
@@ -25,15 +25,14 @@ def test_files_images():
     # test files creation in image classes
     arr = np.zeros((2,3,4))
     aff = np.eye(4)
-    for img_def in class_map.values():
-        klass = img_def['class']
+    for klass in all_image_classes:
         file_map = klass.make_file_map()
         for key, value in file_map.items():
             assert_equal(value.filename, None)
             assert_equal(value.fileobj, None)
             assert_equal(value.pos, 0)
         # If we can't create new images in memory without loading, bail here
-        if not img_def['makeable']:
+        if not klass.makeable:
             continue
         # MGHImage accepts only a few datatypes
         # so we force a type change to float32
@@ -83,22 +82,21 @@ def test_files_interface():
 
 
 def test_round_trip():
-   # write an image to files
-   data = np.arange(24, dtype='i4').reshape((2,3,4))
-   aff = np.eye(4)
-   klasses = [val['class'] for key, val in class_map.items()
-              if val['rw']]
-   for klass in klasses:
-       file_map = klass.make_file_map()
-       for key in file_map:
-           file_map[key].fileobj = BytesIO()
-       img = klass(data, aff)
-       img.file_map = file_map
-       img.to_file_map()
-       # read it back again from the written files
-       img2 = klass.from_file_map(file_map)
-       assert_array_equal(img2.get_data(), data)
-       # write, read it again
-       img2.to_file_map()
-       img3 = klass.from_file_map(file_map)
-       assert_array_equal(img3.get_data(), data)
+    # write an image to files
+    data = np.arange(24, dtype='i4').reshape((2,3,4))
+    aff = np.eye(4)
+    klasses = filter(lambda klass: klass.rw, all_image_classes)
+    for klass in klasses:
+        file_map = klass.make_file_map()
+        for key in file_map:
+            file_map[key].fileobj = BytesIO()
+        img = klass(data, aff)
+        img.file_map = file_map
+        img.to_file_map()
+        # read it back again from the written files
+        img2 = klass.from_file_map(file_map)
+        assert_array_equal(img2.get_data(), data)
+        # write, read it again
+        img2.to_file_map()
+        img3 = klass.from_file_map(file_map)
+        assert_array_equal(img3.get_data(), data)

nibabel/tests/test_image_load_save.py

@@ -26,7 +26,7 @@
 from .. import loadsave as nils
 from .. import (Nifti1Image, Nifti1Header, Nifti1Pair, Nifti2Image, Nifti2Pair,
                 Minc1Image, Minc2Image, Spm2AnalyzeImage, Spm99AnalyzeImage,
-                AnalyzeImage, MGHImage, class_map)
+                AnalyzeImage, MGHImage, all_image_classes)
 
 from ..tmpdirs import InTemporaryDirectory
 
@@ -53,16 +53,14 @@ def test_conversion():
     affine = np.diag([1, 2, 3, 1])
     for npt in np.float32, np.int16:
         data = np.arange(np.prod(shape), dtype=npt).reshape(shape)
-        for r_class_def in class_map.values():
-            r_class = r_class_def['class']
-            if not r_class_def['makeable']:
+        for r_class in all_image_classes:
+            if not r_class.makeable:
                 continue
             img = r_class(data, affine)
             img.set_data_dtype(npt)
-            for w_class_def in class_map.values():
-                if not w_class_def['makeable']:
+            for w_class in all_image_classes:
+                if not w_class.makeable:
                     continue
-                w_class = w_class_def['class']
                 img2 = w_class.from_image(img)
                 assert_array_equal(img2.get_data(), data)
                 assert_array_equal(img2.affine, affine)