Merge pull request #603 from astrofrog/fix-stokes-io

Fix bug when using StokesSpectralCube.read

spectral_cube/io/casa_image.py

@@ -32,6 +32,12 @@
 
 
 def is_casa_image(origin, filepath, fileobj, *args, **kwargs):
+
+    # See note before StringWrapper definition
+    from .core import StringWrapper
+    if len(args) > 0 and isinstance(args[0], StringWrapper):
+        filepath = args[0].value
+
     return filepath is not None and filepath.lower().endswith('.image')
 
 
@@ -181,6 +187,10 @@ def load_casa_image(filename, skipdata=False,
     memory.
     """
 
+    from .core import StringWrapper
+    if isinstance(filename, StringWrapper):
+        filename = filename.value
+
     try:
         import casatools
         iatool = casatools.image
@@ -349,3 +359,7 @@ def load_casa_image(filename, skipdata=False,
 io_registry.register_reader('casa', BaseSpectralCube, load_casa_image)
 io_registry.register_reader('casa_image', BaseSpectralCube, load_casa_image)
 io_registry.register_identifier('casa', BaseSpectralCube, is_casa_image)
+
+io_registry.register_reader('casa', StokesSpectralCube, load_casa_image)
+io_registry.register_reader('casa_image', StokesSpectralCube, load_casa_image)
+io_registry.register_identifier('casa', StokesSpectralCube, is_casa_image)

spectral_cube/io/core.py

@@ -90,6 +90,15 @@
 """
 
 
+# Due to a bug in the astropy I/O infrastructure which causes an exception
+# for directories (which we need for .image), we need to wrap the filenames in
+# a custom string so that astropy doesn't try and call get_readable_fileobj on
+# them.
+class StringWrapper:
+    def __init__(self, value):
+        self.value = value
+
+
 class SpectralCubeRead(registry.UnifiedReadWrite):
 
     __doc__ = DOCSTRING_READ_TEMPLATE.format(clsname='SpectralCube',
@@ -105,7 +114,10 @@ def __call__(self, filename, *args, **kwargs):
         if isinstance(filename, PosixPath):
             filename = str(filename)
         kwargs['target_cls'] = BaseSpectralCube
-        return registry.read(BaseSpectralCube, filename, *args, **kwargs)
+        try:
+            return registry.read(BaseSpectralCube, filename, *args, **kwargs)
+        except IsADirectoryError:  # See note above StringWrapper
+            return registry.read(BaseSpectralCube, StringWrapper(filename), *args, **kwargs)
 
 
 class SpectralCubeWrite(registry.UnifiedReadWrite):
@@ -134,7 +146,10 @@ def __call__(self, filename, *args, **kwargs):
         if isinstance(filename, PosixPath):
             filename = str(filename)
         kwargs['target_cls'] = StokesSpectralCube
-        return registry.read(StokesSpectralCube, filename, *args, **kwargs)
+        try:
+            return registry.read(StokesSpectralCube, filename, *args, **kwargs)
+        except IsADirectoryError:  # See note above StringWrapper
+            return registry.read(StokesSpectralCube, StringWrapper(filename), *args, **kwargs)
 
 
 class StokesSpectralCubeWrite(registry.UnifiedReadWrite):
@@ -174,6 +189,6 @@ def normalize_cube_stokes(cube, target_cls=None):
             raise ValueError("Spectral cube is a Stokes cube that "
                             "does not have an I component")
     elif target_cls is StokesSpectralCube and isinstance(cube, BaseSpectralCube):
-        cube = StokesSpectralCube({'I': cube})
+        return StokesSpectralCube({'I': cube})
     else:
         return cube

spectral_cube/tests/test_casafuncs.py

@@ -9,7 +9,7 @@
 
 from ..io.casa_masks import make_casa_mask
 from ..io.casa_image import wcs_casa2astropy
-from .. import SpectralCube, BooleanArrayMask, VaryingResolutionSpectralCube
+from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, VaryingResolutionSpectralCube
 from . import path
 
 try:
@@ -78,12 +78,24 @@ def test_casa_read(filename, tmp_path):
 
     make_casa_testimage(filename, tmp_path / 'casa.image')
 
-    casacube = SpectralCube.read(tmp_path / 'casa.image', format='casa_image')
+    casacube = SpectralCube.read(tmp_path / 'casa.image')
 
     assert casacube.shape == cube.shape
     # what other equalities should we check?
 
 
+@pytest.mark.skipif(not casaOK, reason='CASA tests must be run in a CASA environment.')
+def test_casa_read_stokes(data_advs, tmp_path):
+
+    cube = StokesSpectralCube.read(data_advs)
+
+    make_casa_testimage(data_advs, tmp_path / 'casa.image')
+
+    casacube = StokesSpectralCube.read(tmp_path / 'casa.image')
+
+    assert casacube.I.shape == cube.I.shape
+
+
 @pytest.mark.skipif(not casaOK, reason='CASA tests must be run in a CASA environment.')
 def test_casa_mask(data_adv, tmp_path):
 

spectral_cube/tests/test_io.py

@@ -44,6 +44,15 @@ def test_4d_stokes(data_advs):
     f.close()
 
 
+def test_4d_stokes_read_3d(data_adv):
+    # Regression test for a bug that caused StokesSpectralCube.read to not work
+    # correctly when reading in a 3D FITS file.
+    f = pyfits.open(data_adv)
+    c = StokesSpectralCube.read(f)
+    assert isinstance(c, StokesSpectralCube)
+    f.close()
+
+
 def test_3d_beams(data_vda_beams):
     c = SpectralCube.read(data_vda_beams)
     np.testing.assert_almost_equal(c.beams[0].major.value, 0.4)