Merge 9d91437 into 3205324

CHANGES.rst

@@ -64,6 +64,9 @@ New Features
     arrays are fixed-width and silently drop characters which do not
     fit within the fixed width. [#5624]
 
+  - Added functionality to allow ``astropy.units.Quantity`` to be written 
+    as a normal column to FITS files. [#5910]
+
 - ``astropy.time``
 
 - ``astropy.units``

astropy/io/fits/convenience.py

@@ -72,6 +72,7 @@
 from .fitsrec import FITS_rec
 from ...units import Unit
 from ...units.format.fits import UnitScaleError
+from ...units import Quantity
 from ...extern import six
 from ...extern.six import string_types
 from ...utils.exceptions import AstropyUserWarning
@@ -460,12 +461,17 @@ def table_to_hdu(table):
     # Avoid circular imports
     from .connect import is_column_keyword, REMOVE_KEYWORDS
 
-    # Tables with mixin columns are not supported
+    # Not all tables with mixin columns are supported
     if table.has_mixin_columns:
-        mixin_names = [name for name, col in table.columns.items()
-                       if not isinstance(col, table.ColumnClass)]
-        raise ValueError('cannot write table with mixin column(s) {0}'
-                         .format(mixin_names))
+        #the import is done here, in order to avoid it at build time as erfa is not yet available then
+        from ...table.column import BaseColumn
+
+        # Only those columns which are instances of BaseColumn or Quantity can be written
+        unsupported_cols = table.columns.not_isinstance((BaseColumn, Quantity))
+        if unsupported_cols:
+            unsupported_names = [col.info.name for col in unsupported_cols]
+            raise ValueError('cannot write table with mixin column(s) {0}'
+                         .format(unsupported_names))
 
     # Create a new HDU object
     if table.masked:

astropy/io/fits/tests/test_connect.py

@@ -10,7 +10,7 @@
 
 from .... import units as u
 from ....extern.six.moves import range, zip
-from ....table import Table
+from ....table import Table, QTable
 from ....tests.helper import pytest, catch_warnings
 from ....units.format.fits import UnitScaleError
 
@@ -92,13 +92,14 @@ def test_simple_noextension(self, tmpdir):
         t2 = Table.read(filename)
         assert equal_data(t1, t2)
 
-    def test_with_units(self, tmpdir):
+    @pytest.mark.parametrize('table_type', (Table, QTable))
+    def test_with_units(self, table_type, tmpdir):
         filename = str(tmpdir.join('test_with_units.fits'))
-        t1 = Table(self.data)
+        t1 = table_type(self.data)
         t1['a'].unit = u.m
         t1['c'].unit = u.km / u.s
         t1.write(filename, overwrite=True)
-        t2 = Table.read(filename)
+        t2 = table_type.read(filename)
         assert equal_data(t1, t2)
         assert t2['a'].unit == u.m
         assert t2['c'].unit == u.km / u.s

astropy/table/table.py

@@ -150,6 +150,39 @@ def keys(self):
     def values(self):
         return list(OrderedDict.values(self))
 
+    def isinstance(self, cls):
+        """
+        Return a list of columns which are instances of the specified classes.
+
+        Parameters
+        ----------
+        cls : class or tuple of classes
+            Column class (including mixin) or tuple of Column classes.
+
+        Returns
+        -------
+        col_list : list of Columns
+            List of Column objects which are instances of given classes.
+        """
+        cols = [col for col in self.values() if isinstance(col, cls)]
+        return cols
+
+    def not_isinstance(self, cls):
+        """
+        Return a list of columns which are not instances of the specified classes.
+
+        Parameters
+        ----------
+        cls : class or tuple of classes
+            Column class (including mixin) or tuple of Column classes.
+
+        Returns
+        -------
+        col_list : list of Columns
+            List of Column objects which are not instances of given classes.
+        """
+        cols = [col for col in self.values() if not isinstance(col, cls)]
+        return cols
 
 class Table(object):
     """A class to represent tables of heterogeneous data.
@@ -853,7 +886,7 @@ def __bytes__(self):
     def has_mixin_columns(self):
         """
         True if table has any mixin columns (defined as columns that are not Column
-        subclasses)
+        subclasses).
         """
         return any(has_info_class(col, MixinInfo) for col in self.columns.values())
 

astropy/table/tests/test_mixin.py

@@ -24,6 +24,7 @@
 from ... import time
 from ... import coordinates
 from ... import units as u
+from ..column import BaseColumn
 from .. import table_helpers
 from .conftest import MIXIN_COLS
 
@@ -102,14 +103,44 @@ def test_io_ascii_write():
             t.write(out, format=fmt['Format'])
 
 
+def test_io_quantity_write(tmpdir):
+    """
+    Test that table with Quantity mixin column can be written by io.fits,
+    but not by io.votable and io.misc.hdf5. Validation of the output is done.
+    Test that io.fits writes a table containing Quantity mixin columns that can
+    be round-tripped (metadata unit).
+    """
+    t = QTable()
+    t['a'] = u.Quantity([1,2,4], unit='Angstrom')
+
+    filename = tmpdir.join("table-tmp").strpath
+    open(filename, 'w').close()
+
+    for fmt in ('fits', 'votable', 'hdf5'):
+        if fmt == 'fits':
+            t.write(filename, format=fmt, overwrite=True)
+            qt = QTable.read(filename, format=fmt)
+            assert isinstance(qt['a'], u.Quantity)
+            assert qt['a'].unit == 'Angstrom'
+            continue
+        if fmt == 'hdf5' and not HAS_H5PY:
+            continue
+        with pytest.raises(ValueError) as err:
+            t.write(filename, format=fmt, overwrite=True)
+        assert 'cannot write table with mixin column(s)' in str(err.value)
+
+
 def test_io_write_fail(mixin_cols):
     """
-    Test that table with mixin column cannot be written by io.votable,
-    io.fits, and io.misc.hdf5
-    every pure Python writer.  No validation of the output is done,
-    this just confirms no exceptions.
+    Test that table with mixin column (excluding Quantity) cannot be written by io.votable,
+    io.fits, and io.misc.hdf5.
     """
     t = QTable(mixin_cols)
+    #Only do this test if there are unsupported column types (i.e. anything besides
+    #BaseColumn or Quantity subclasses.
+    unsupported_cols = t.columns.not_isinstance((BaseColumn, u.Quantity))
+    if not unsupported_cols:
+        pytest.skip("no unsupported column types")
     for fmt in ('fits', 'votable', 'hdf5'):
         if fmt == 'hdf5' and not HAS_H5PY:
             continue