Merge e339e65 into 3037a6f

.github/workflows/python-package.yml

@@ -89,7 +89,7 @@ jobs:
               with:
                 python-version: ${{ matrix.python-version }}
             - name: Install Python dependencies
-              run: python -m pip install --upgrade pip setuptools wheel Sphinx sphinx-rtd-theme
+              run: python -m pip install --upgrade pip setuptools wheel Sphinx\<7 sphinx-rtd-theme
             - name: Test the documentation
               run: sphinx-build -W --keep-going -b html doc doc/_build/html
 

.readthedocs.yml

@@ -5,16 +5,20 @@
 # Required
 version: 2
 
+build:
+    os: "ubuntu-22.04"
+    tools:
+        python: "3.10"
+
 # Build documentation in the doc/ directory with Sphinx
 sphinx:
     configuration: doc/conf.py
-    fail_on_warning: true
 
 # Optionally build your docs in additional formats such as PDF and ePub
-formats: all
+# formats: all
 
 # Optionally set the version of Python and requirements required to build your docs
 python:
-    version: "3.8"
     install:
-      - requirements: doc/rtd-requirements.txt
+        - requirements: doc/rtd-requirements.txt
+

doc/api.rst

@@ -8,8 +8,8 @@ desiutil API
 .. automodule:: desiutil.annotate
     :members:
 
-.. .. automodule:: desiutil.api
-..     :members:
+.. automodule:: desiutil.api
+    :members:
 
 .. automodule:: desiutil.bitmask
     :members:
@@ -53,8 +53,8 @@ desiutil API
 .. automodule:: desiutil.redirect
     :members:
 
-.. .. automodule:: desiutil.setup
-..     :members:
+.. automodule:: desiutil.setup
+    :members:
 
 .. automodule:: desiutil.sklearn
     :members:

doc/changes.rst

@@ -9,6 +9,14 @@ Change Log
 
 * Remove deprecated commands in :mod:`desiutil.setup`.
 * Remove deprecated module :mod:`desiutil.census`.
+* Remove deprecated top-level ``setup.py``.
+
+3.4.1 (unreleased)
+------------------
+
+* Fully support adding units and comments to FITS table columns (PR `#201`).
+
+.. _`#201`: https://github.com/desihub/desiutil/pull/201
 
 3.4.0 (2023-09-14)
 ------------------

py/desiutil/annotate.py

@@ -11,14 +11,21 @@
 import os
 import sys
 from argparse import ArgumentParser
+from warnings import warn
 import yaml
 from astropy.io import fits
 from astropy.table import Table, QTable
-from astropy.units import UnitConversionError
+from astropy.units import Unit, UnitConversionError
 from . import __version__ as desiutilVersion
 from .log import get_logger, DEBUG
 
 
+class FITSUnitWarning(UserWarning):
+    """Warnings related to invalid FITS units.
+    """
+    pass
+
+
 def find_column_name(columns, prefix=('unit', )):
     """Find the column that contains unit descriptions, or comments.
 
@@ -74,6 +81,90 @@ def find_key_name(data, prefix=('unit', )):
     raise KeyError(f"No key matching '{prefix[0]}' found!")
 
 
+def validate_unit(unit, error=False):
+    """Check units for consistency with FITS standard, while allowing
+    some special exceptions.
+
+    Parameters
+    ----------
+    unit : :class:`str`
+        The unit to parse.
+    error : :class:`bool`, optional
+        If ``True``, failure to interpret the unit raises an
+        exception.
+
+    Returns
+    -------
+    :class:`str`
+        If a special exception is detected, the name of the unit
+        is returned.  Otherwise, ``None``.
+
+    Raises
+    ------
+    :exc:`ValueError`
+        If `error` is set and `unit` can't be parsed.
+    """
+    if unit is None:
+        return None
+    acceptable_units = ('maggie', 'maggy', 'mgy',
+                        'electron/Angstrom',
+                        'log(Angstrom)')
+    try:
+        au = Unit(unit, format='fits')
+    except ValueError as e:
+        m = str(e)
+        bad_unit = m.split()[0]
+        if any([u in bad_unit for u in acceptable_units]) and 'Numeric factor' not in m:
+            return bad_unit
+        else:
+            if error:
+                raise
+            else:
+                warn(m, FITSUnitWarning)
+    return None
+
+
+def check_comment_length(comments, error=True):
+    """Ensure keyword comments are short enough that they will not be truncated.
+
+    By experiment, :mod:`astropy.io.fits` will truncate comments longer than
+    46 characters, however the warning it emits when it does so is not very
+    informative.
+
+    Parameters
+    ----------
+    comments : :class:`dict`
+        Mapping of table columns to comments.
+    error : :class:`bool`, optional
+        If ``False`` just warn about long comments instead of raising an exception.
+
+    Returns
+    -------
+    :class:`int`
+        The number of long comments detected, although the value is only relevant
+        if `error` is ``False``.
+
+    Raises
+    ------
+    ValueError
+        If any comment is too long.
+    """
+    log = get_logger()
+    too_long = 46
+    n_long = 0
+    for key in comments:
+        if len(comments[key]) > too_long:
+            n_long += 1
+            if error:
+                log.error("Long comment detected for '%s'!", key)
+            else:
+                log.warning("Long comment detected for '%s', will be truncated to '%s'!",
+                            key, comments[key][:too_long])
+    if n_long > 0 and error:
+        raise ValueError(f"{n_long:d} long comments detected!")
+    return n_long
+
+
 def load_csv_units(filename):
     """Parse a CSV file that contains column names and units and optionally comments.
 
@@ -228,8 +319,11 @@ def annotate_table(table, units, inplace=False):
     return t
 
 
-def annotate(filename, extension, units=None, comments=None):
-    """Add annotations to `filename`.
+def annotate_fits(filename, extension, output, units=None, comments=None,
+                  truncate=False, overwrite=False):
+    """Add annotations to HDU `extension` in FITS file `filename`.
+
+    HDU `extension` must be a :class:`astropy.io.fits.BinTableHDU`.
 
     If `units` or `comments` is an empty dictionary, it will be ignored.
 
@@ -239,64 +333,82 @@ def annotate(filename, extension, units=None, comments=None):
         Name of FITS file.
     extension : :class:`str` or :class:`int`
         Name or number of extension in `filename`.
+    output : :class:`str`
+        Name of file to write to.
     units : :class:`dict`, optional
         Mapping of table columns to units.
     comments : :class:`dict`, optional
         Mapping of table columns to comments.
+    truncate : :class:`bool`, optional
+        Allow long comments to be truncated when written out. The default
+        is to raise an error if a comment is too long.
+    overwrite : :class:`bool`, optional
+        Pass this keyword to :meth:`astropy.io.fits.HDUList.writeto`.
+
 
     Returns
     -------
     :class:`astropy.io.fits.HDUList`
-        An updated version of the file.
+        An updated version of the file, equivalent to the data in `output`.
+
+    Raises
+    ------
+    IndexError
+        If the HDU specified (numerically) by `extension` does not exist.
+    KeyError
+        If the HDU specified (as an ``EXTNAME``) by `extension` does not exist.
+    TypeError
+        If the HDU specified is not supported by this function.
+    ValueError
+        If neither `units` nor `comments` are specified.
+
+    Notes
+    -----
+    * Due to the way :func:`astropy.io.fits.open` manages memory, changes
+      have to be written out while `filename` is still open,
+      hence the mandatory `output` argument.
+    * A FITS HDU checksum will always be added to the output, even if it
+      was not already present.
     """
     log = get_logger()
-    new_hdus = list()
-    with fits.open(filename, mode='readonly', memmap=False, lazy_load_hdus=False, uint=False, disable_image_compression=True, do_not_scale_image_data=True, character_as_bytes=True, scale_back=True) as hdulist:
-        log.debug(hdulist._open_kwargs)
-        kwargs = hdulist._open_kwargs.copy()
-        for h in hdulist:
-            hc = h.copy()
-            if hasattr(h, '_do_not_scale_image_data'):
-                hc._do_not_scale_image_data = h._do_not_scale_image_data
-            if hasattr(h, '_bzero'):
-                hc._bzero = h._bzero
-            if hasattr(h, '_bscale'):
-                hc._bzero = h._bscale
-            if hasattr(h, '_scale_back'):
-                hc._scale_back = h._scale_back
-            if hasattr(h, '_uint'):
-                hc._uint = h._uint
-            #
-            # Work around header comments not copied for BinTableHDU.
-            #
-            if isinstance(h, fits.BinTableHDU):
-                for key in h.header.keys():
-                    hc.header.comments[key] = h.header.comments[key]
-            #
-            # Work around disappearing BZERO and BSCALE keywords.
-            #
-            if isinstance(h, fits.ImageHDU) and 'BZERO' in h.header and 'BSCALE' in h.header:
-                if 'BZERO' not in hc.header or 'BSCALE' not in hc.header:
-                    iscale = h.header.index('BSCALE')
-                    izero = h.header.index('BZERO')
-                    if izero > iscale:
-                        hc.header.insert(iscale - 1, ('BSCALE', h.header['BSCALE'], h.header.comments['BSCALE']), after=True)
-                        hc.header.insert(iscale, ('BZERO', h.header['BZERO'], h.header.comments['BZERO']), after=True)
-                    else:
-                        hc.header.insert(izero - 1, ('BZERO', h.header['BZERO'], h.header.comments['BZERO']), after=True)
-                        hc.header.insert(izero, ('BSCALE', h.header['BSCALE'], h.header.comments['BSCALE']), after=True)
-            new_hdus.append(hc)
-    new_hdulist = fits.HDUList(new_hdus)
-    new_hdulist._open_kwargs = kwargs
-    log.debug(new_hdulist._open_kwargs)
     try:
         ext = int(extension)
     except ValueError:
         ext = extension
-    try:
-        hdu = new_hdulist[ext]
-    except (IndexError, KeyError):
-        raise
+    if not units and not comments:
+        raise ValueError("No input units or comments specified!")
+    if comments:
+        n_long = check_comment_length(comments, error=(not truncate))
+    with fits.open(filename, mode='readonly') as hdulist:
+        new_hdulist = hdulist.copy()
+        try:
+            hdu = new_hdulist[ext]
+        except (IndexError, KeyError):
+            raise
+        if isinstance(hdu, fits.BinTableHDU) and not isinstance(hdu, fits.CompImageHDU):
+            #
+            # fits.CompImageHDU is a subclass of fits.BinTableHDU.
+            #
+            for i in range(1, 1000):
+                ttype = f"TTYPE{i:d}"
+                if ttype not in hdu.header:
+                    break
+                colname = hdu.header[ttype]
+                if comments and colname in comments and comments[colname].strip():
+                    if hdu.header.comments[ttype].strip():
+                        log.warning("Overriding comment on column '%s': '%s' -> '%s'.", colname, hdu.header.comments[ttype].strip(), comments[colname].strip())
+                    hdu.header[ttype] = (colname, comments[colname].strip())
+                if units and colname in units and units[colname].strip():
+                    tunit = f"TUNIT{i:d}"
+                    if tunit in hdu.header and hdu.header[tunit].strip():
+                        log.warning("Overriding units for column '%s': '%s' -> '%s'.", colname, hdu.header[tunit].strip(), units[colname].strip())
+                        hdu.header[tunit] = (units[colname].strip(), colname+' units')
+                    else:
+                        hdu.header.insert(f"TFORM{i:d}", (tunit, units[colname].strip(), colname+' units'), after=True)
+        else:
+            raise TypeError("Adding units to objects other than fits.BinTableHDU is not supported!")
+        hdu.add_checksum()
+        new_hdulist.writeto(output, output_verify='warn', overwrite=overwrite, checksum=False)
     return new_hdulist
 
 
@@ -309,12 +421,14 @@ def _options():
                         help="COMMENTS should have the form COLUMN='comment':COLUMN='comment'.")
     parser.add_argument('-C', '--csv', action='store', dest='csv', metavar='CSV',
                         help="Read annotations from CSV file.")
+    parser.add_argument('-D', '--disable-comments', action='store_true', dest='disable_comments',
+                        help='Do not add comments, even if they are defined in one of the inputs.')
     parser.add_argument('-e', '--extension', dest='extension', action='store', metavar='EXT', default='1',
                         help="Update FITS extension EXT, which can be a number or an EXTNAME. If not specified, HDU 1 will be updated, which is standard for simple binary tables.")
     parser.add_argument('-o', '--overwrite', dest='overwrite', action='store_true',
                         help='Overwrite the input FITS file.')
-    parser.add_argument('-t', '--test', dest='test', action='store_true',
-                        help='Test mode; show what would be done but do not change any files.')
+    parser.add_argument('-T', '--truncate-comments', dest='truncate', action='store_true',
+                        help='Allow any long comments to be truncated when written out.')
     parser.add_argument('-u', '--units', action='store', dest='units', metavar='UNITS',
                         help="UNITS should have the form COLUMN='unit':COLUMN='unit'.")
     parser.add_argument('-v', '--verbose', dest='verbose', action='store_true',
@@ -338,7 +452,7 @@ def main():
         An integer suitable for passing to :func:`sys.exit`.
     """
     options = _options()
-    if options.test or options.verbose:
+    if options.verbose:
         log = get_logger(DEBUG)
     else:
         log = get_logger()
@@ -357,24 +471,37 @@ def main():
         else:
             log.debug("No comments have been specified.")
             comments = dict()
-    log.debug("units = %s", units)
-    log.debug("comments = %s", comments)
-    hdulist = annotate(options.fits, options.extension, units, comments)
+    if options.disable_comments:
+        log.info("Comments are disabled by user request.")
+        comments = dict()
+    if units:
+        log.debug("Input Units")
+        for k, v in units.items():
+            log.debug("'%s': '%s'", k, v)
+    if comments:
+        log.debug("Input Comments")
+        for k, v in comments.items():
+            log.debug("'%s': '%s'", k, v)
     if options.overwrite and options.output:
         output = options.output
     elif options.overwrite:
         output = options.fits
     elif options.output:
         output = options.output
     else:
-        log.error("--overwrite not specified and no output file specified!")
+        log.critical("--overwrite not specified and no output file specified!")
         return 1
     try:
-        hdulist.writeto(output, output_verify='warn', overwrite=options.overwrite, checksum=False)
+        hdulist = annotate_fits(options.fits, options.extension, output,
+                                units, comments, truncate=options.truncate,
+                                overwrite=options.overwrite)
     except OSError as e:
         if 'overwrite' in e.args[0]:
-            log.error("Output file exists and --overwrite was not specified!")
+            log.critical("Output file exists and --overwrite was not specified!")
         else:
-            log.error(e.args[0])
+            log.critical(e.args[0])
+        return 1
+    except (IndexError, KeyError, TypeError, ValueError) as e:
+        log.critical(str(e))
         return 1
     return 0