BF - make nifti1 load and save allow single or pair-type filenames

Author: matthew-brett

nibabel/nifti1.py

@@ -12,7 +12,7 @@
 import numpy.linalg as npl
 
 from .volumeutils import Recoder, make_dt_codes, endian_codes
-from .spatialimages import HeaderDataError
+from .spatialimages import HeaderDataError, ImageFileError
 from .batteryrunners import Report
 from .quaternions import fillpositive, quat2mat, mat2quat
 from . import analyze # module import
@@ -1404,5 +1404,41 @@ def update_header(self):
             hdr['vox_offset'] = min_vox_offset
 
 
-load = Nifti1Image.load
-save = Nifti1Image.instance_to_filename
+def load(filename):
+    """ Load nifti1 single or pair from `filename`
+
+    Parameters
+    ----------
+    filename : str
+        filename of image to be loaded
+
+    Returns
+    -------
+    img : Nifti1Image or Nifti1Pair
+        nifti1 single or pair image instance
+
+    Raises
+    ------
+    ImageFileError: if `filename` doesn't look like nifti1
+    IOError : if `filename` does not exist
+    """
+    try:
+        img = Nifti1Image.load(filename)
+    except ImageFileError:
+        return Nifti1Pair.load(filename)
+    return img
+
+
+def save(img, filename):
+    """ Save nifti1 single or pair to `filename`
+
+    Parameters
+    ----------
+    filename : str
+        filename to which to save image
+    """
+    try:
+        Nifti1Image.instance_to_filename(img, filename)
+    except ImageFileError:
+        Nifti1Pair.instance_to_filename(img, filename)
+

nibabel/tests/test_nifti1.py

@@ -26,6 +26,7 @@
                         assert_raises)
 from nose import SkipTest
 
+from ..tmpdirs import InTemporaryDirectory
 from ..testing import parametric, data_path
 
 from . import test_analyze as tana
@@ -531,3 +532,18 @@ def test_recoded_fields():
     yield assert_equal(hdr.get_value_label('slice_code'),
                        'alternating decreasing')
 
+
+def test_load():
+    # test module level load.  We try to load a nii and an .img and a .hdr and
+    # expect to get a nifti back of single or pair type
+    arr = np.arange(24).reshape((2,3,4))
+    aff = np.diag([2, 3, 4, 1])
+    simg = Nifti1Image(arr, aff)
+    pimg = Nifti1Pair(arr, aff)
+    with InTemporaryDirectory():
+        nifti1.save(simg, 'test.nii')
+        assert_array_equal(arr, nifti1.load('test.nii').get_data())
+        nifti1.save(simg, 'test.img')
+        assert_array_equal(arr, nifti1.load('test.img').get_data())
+        nifti1.save(simg, 'test.hdr')
+        assert_array_equal(arr, nifti1.load('test.hdr').get_data())