Merge pull request #600 from astrofrog/unified-io

Update I/O infrastructure to use Astropy unified I/O

CHANGES.rst

@@ -6,6 +6,8 @@
 
 - Refactor package infrastructure to no longer use astropy-helpers. #599
 
+- Switch to using unified I/O infrastructure from Astropy. #600
+
 0.4.5 (unreleased)
 ------------------
  - Added support for casatools-based io in #541 and beam reading from CASA

spectral_cube/__init__.py

@@ -11,6 +11,14 @@
                     FunctionMask)
 from .lower_dimensional_structures import (OneDSpectrum, Projection, Slice)
 
+# Import the following sub-packages to make sure the I/O functions are registered
+from .io import casa_image
+del casa_image
+from .io import class_lmv
+del class_lmv
+from .io import fits
+del fits
+
 __all__ = ['SpectralCube', 'VaryingResolutionSpectralCube',
             'StokesSpectralCube', 'CompositeMask', 'LazyComparisonMask',
             'LazyMask', 'BooleanArrayMask', 'FunctionMask',

spectral_cube/io/casa_image.py

@@ -9,14 +9,16 @@
 from astropy import units as u
 from astropy.wcs.wcsapi.sliced_low_level_wcs import sanitize_slices
 from astropy import log
+from astropy.io import registry as io_registry
 import numpy as np
 from radio_beam import Beam, Beams
 
 import dask.array
 
 from .. import SpectralCube, StokesSpectralCube, BooleanArrayMask, LazyMask, VaryingResolutionSpectralCube
+from ..spectral_cube import BaseSpectralCube
 from .. import cube_utils
-from .. utils import BeamWarning, cached
+from .. utils import BeamWarning, cached, StokesWarning
 from .. import wcs_utils
 
 # Read and write from a CASA image. This has a few
@@ -29,11 +31,8 @@
 # yield the same array in memory that we would get from astropy.
 
 
-def is_casa_image(input, **kwargs):
-    if isinstance(input, six.string_types):
-        if input.endswith('.image'):
-            return True
-    return False
+def is_casa_image(origin, filepath, fileobj, *args, **kwargs):
+    return filepath is not None and filepath.lower().endswith('.image')
 
 
 def wcs_casa2astropy(ia, coordsys):
@@ -174,7 +173,7 @@ def __getitem__(self, value):
 
 
 def load_casa_image(filename, skipdata=False,
-                    skipvalid=False, skipcs=False, **kwargs):
+                    skipvalid=False, skipcs=False, target_cls=None, **kwargs):
     """
     Load a cube (into memory?) from a CASA image. By default it will transpose
     the cube into a 'python' order and drop degenerate axes. These options can
@@ -343,5 +342,10 @@ def load_casa_image(filename, skipdata=False,
         raise ValueError("CASA image has {0} dimensions, and therefore "
                          "is not readable by spectral-cube.".format(wcs.naxis))
 
+    from .core import normalize_cube_stokes
+    return normalize_cube_stokes(cube, target_cls=target_cls)
 
-    return cube
+
+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)

spectral_cube/io/class_lmv.py

@@ -6,6 +6,8 @@
 import warnings
 import string
 from astropy import log
+from astropy.io import registry as io_registry
+from ..spectral_cube import BaseSpectralCube
 from .fits import load_fits_cube
 
 """
@@ -40,17 +42,13 @@
               7:'GLS', 8:'SFL', }
 _bunit_dict = {'k (tmb)': 'K'}
 
-def is_lmv(input, **kwargs):
+def is_lmv(origin, filepath, fileobj, *args, **kwargs):
     """
     Determine whether input is in GILDAS CLASS lmv format
     """
-    if isinstance(input, six.string_types):
-        if input.lower().endswith(('.lmv')):
-            return True
-    else:
-        return False
+    return filepath is not None and filepath.lower().endswith('.lmv')
 
-def read_lmv(filename):
+def read_lmv(lf):
     """
     Read an LMV cube file
 
@@ -59,29 +57,28 @@ def read_lmv(filename):
     log.warning("CLASS LMV cube reading is tentatively supported.  "
              "Please post bug reports at the first sign of danger!")
 
-    with open(filename,'rb') as lf:
-        # lf for "LMV File"
-        filetype = _read_string(lf, 12)
-        #!---------------------------------------------------------------------
-        #! @ private
-        #!       SYCODE system code
-        #!       '-'    IEEE
-        #!       '.'    EEEI (IBM like)
-        #!       '_'    VAX
-        #!       IMCODE file code
-        #!       '<'    IEEE  64 bits    (Little Endian, 99.9 % of recent computers)
-        #!       '>'    EEEI  64 bits    (Big Endian, HPUX, IBM-RISC, and SPARC ...)
-        #!---------------------------------------------------------------------
-        imcode = filetype[6]
-        if filetype[:6] != 'GILDAS' or filetype[7:] != 'IMAGE':
-            raise TypeError("File is not a GILDAS Image file")
-
-        if imcode in ('<','>'):
-            if imcode =='>':
-                log.warning("Swap the endianness first...")
-            return read_lmv_type2(lf)
-        else:
-            return read_lmv_type1(lf)
+    # lf for "LMV File"
+    filetype = _read_string(lf, 12)
+    #!---------------------------------------------------------------------
+    #! @ private
+    #!       SYCODE system code
+    #!       '-'    IEEE
+    #!       '.'    EEEI (IBM like)
+    #!       '_'    VAX
+    #!       IMCODE file code
+    #!       '<'    IEEE  64 bits    (Little Endian, 99.9 % of recent computers)
+    #!       '>'    EEEI  64 bits    (Big Endian, HPUX, IBM-RISC, and SPARC ...)
+    #!---------------------------------------------------------------------
+    imcode = filetype[6]
+    if filetype[:6] != 'GILDAS' or filetype[7:] != 'IMAGE':
+        raise TypeError("File is not a GILDAS Image file")
+
+    if imcode in ('<','>'):
+        if imcode =='>':
+            log.warning("Swap the endianness first...")
+        return read_lmv_type2(lf)
+    else:
+        return read_lmv_type1(lf)
 
 def read_lmv_type1(lf):
     header = {}
@@ -244,9 +241,9 @@ def read_lmv_type1(lf):
     # debug
     #return data.reshape([naxis3,naxis2,naxis1]), header, hdr_f, hdr_s, hdr_i, hdr_d, hdr_d_2
 
-def read_lmv_tofits(filename):
+def read_lmv_tofits(fileobj):
     from astropy.io import fits
-    data,header = read_lmv(filename)
+    data,header = read_lmv(fileobj)
     # LMV may contain extra dimensions that are improperly labeled
     data = data.squeeze()
     bad_kws = ['NAXIS4','CRVAL4','CRPIX4','CDELT4','CROTA4','CUNIT4','CTYPE4']
@@ -262,9 +259,9 @@ def read_lmv_tofits(filename):
     hdu = fits.PrimaryHDU(data=data, header=Header)
     return hdu
 
-def load_lmv_cube(filename):
-    hdu = read_lmv_tofits(filename)
-    meta = {'filename':filename}
+def load_lmv_cube(fileobj, target_cls=None):
+    hdu = read_lmv_tofits(fileobj)
+    meta = {'filename':fileobj.name}
 
     return load_fits_cube(hdu, meta=meta)
 
@@ -671,3 +668,8 @@ def read_lmv_type2(lf):
     data[data==bval] = np.nan
 
     return data,header
+
+
+io_registry.register_reader('lmv', BaseSpectralCube, load_lmv_cube)
+io_registry.register_reader('class_lmv', BaseSpectralCube, load_lmv_cube)
+io_registry.register_identifier('lmv', BaseSpectralCube, is_lmv)