STWCS update to insure that NEXTEND keyword (if already present in primary header) remains consistent with full size of FITS file after updates.

git-svn-id: http://svn.stsci.edu/svn/ssb/stsci_python/stwcs/trunk@28702 fe389314-cf27-0410-b35b-8c050e845b92

3 files changed, 34 insertions, 10 deletions

diff --git a/lib/stwcs/updatewcs/utils.py b/lib/stwcs/updatewcs/utils.py
index 0f3d6cc..550ce56 100644
--- a/lib/stwcs/updatewcs/utils.py
+++ b/lib/stwcs/updatewcs/utils.py
@@ -29,6 +29,19 @@ def getBinning(fobj, extver=1):
         binned = fobj['SCI', extver].header.get('BINAXIS',1)
     return binned
 
+def updateNEXTENDKw(fobj):
+    """
+    Updates PRIMARY header with correct value for NEXTEND, if present.
+
+    Parameters
+    -----------
+    fobj : pyfits.HDUList
+        PyFITS object for file opened in update mode
+
+    """
+    if 'nextend' in fobj[0].header:
+        fobj[0].header['nextend'] = len(fobj) -1
+
 def extract_rootname(kwvalue,suffix=""):
     """ Returns the rootname from a full reference filename
 
@@ -49,7 +62,7 @@ def extract_rootname(kwvalue,suffix=""):
         fullval = kwvalue
     # Extract filename without path from kwvalue
     fname = os.path.basename(fullval).strip()
-    
+
     # Now, rip out just the rootname from the full filename
     rootname = fileutil.buildNewRootname(fname)
 
@@ -108,7 +121,7 @@ def build_default_wcsname(idctab):
 def build_sipname(fobj, fname=None, sipname=None):
     """
     Build a SIPNAME from IDCTAB
-    
+
     Parameters
     ----------
     fobj: HDUList
@@ -117,7 +130,7 @@ def build_sipname(fobj, fname=None, sipname=None):
           science file name (to be used if ROOTNAMe is not present
     sipname: string
           user supplied SIPNAME keyword
-          
+
     Returns
     -------
     sipname, idctab
@@ -136,7 +149,7 @@ def build_sipname(fobj, fname=None, sipname=None):
             sipname = fobj[0].header["SIPNAME"]
         except KeyError:
             try:
-                idcname = extract_rootname(fobj[0].header["IDCTAB"],suffix='_idc')              
+                idcname = extract_rootname(fobj[0].header["IDCTAB"],suffix='_idc')
                 try:
                     rootname = fobj[0].header['rootname']
                 except KeyError:
@@ -153,14 +166,14 @@ def build_sipname(fobj, fname=None, sipname=None):
 def build_npolname(fobj, npolfile=None):
     """
     Build a NPOLNAME from NPOLFILE
-    
+
     Parameters
     ----------
     fobj: HDUList
           pyfits file object
     npolfile: string
           user supplied NPOLFILE keyword
-          
+
     Returns
     -------
     npolname, npolfile
@@ -186,14 +199,14 @@ def build_npolname(fobj, npolfile=None):
 def build_d2imname(fobj, d2imfile=None):
     """
     Build a D2IMNAME from D2IMFILE
-    
+
     Parameters
     ----------
     fobj: HDUList
           pyfits file object
     d2imfile: string
           user supplied NPOLFILE keyword
-          
+
     Returns
     -------
     d2imname, d2imfile
diff --git a/lib/stwcs/wcsutil/headerlet.py b/lib/stwcs/wcsutil/headerlet.py
index 0cad22a..c2abad0 100644
--- a/lib/stwcs/wcsutil/headerlet.py
+++ b/lib/stwcs/wcsutil/headerlet.py
@@ -193,7 +193,6 @@ def get_header_kw_vals(hdr, kwname, kwval, default=0):
             kwval = default
     return kwval
 
-
 @with_logging
 def find_headerlet_HDUs(fobj, hdrext=None, hdrname=None, distname=None,
                         strict=True, logging=False, logmode='w'):
@@ -776,6 +775,7 @@ def write_headerlet(filename, hdrname, output=None, sciext='SCI',
                 wcscorr.update_wcscorr(fobj, source=hdrletobj,
                                        extname='SIPWCS', wcs_id=wname)
 
+                utils.updateNEXTENDKw(fobj)
                 fobj.flush()
             else:
                 message = """
@@ -1313,6 +1313,7 @@ def _delete_single_headerlet(filename, hdrname=None, hdrext=None, distname=None,
     for hdrind in hdrlet_ind:
         del fobj[hdrind]
 
+    utils.updateNEXTENDKw(fobj)
     # Update file object with changes
     fobj.flush()
     # close file, if was opened by this function
@@ -1506,6 +1507,7 @@ def restore_from_headerlet(filename, hdrname=None, hdrext=None, archive=True,
     #
     hdrlet.apply_as_primary(fobj, attach=False, archive=archive, force=force)
 
+    utils.updateNEXTENDKw(fobj)
     fobj.flush()
     if close_fobj:
         fobj.close()
@@ -1618,6 +1620,7 @@ def restore_all_with_distname(filename, distname, primary, archive=True,
                 hdrlet.apply_as_alternate(fobj, attach=False,
                                           wcsname=hdrlet[0].header['wcsname'])
 
+    utils.updateNEXTENDKw(fobj)
     fobj.flush()
     if close_fobj:
         fobj.close()
@@ -1743,6 +1746,7 @@ def archive_as_headerlet(filename, hdrname, sciext='SCI',
 
         fobj.append(hlt_hdu)
 
+        utils.updateNEXTENDKw(fobj)
         fobj.flush()
     else:
         message = """
@@ -2051,6 +2055,7 @@ class Headerlet(pyfits.HDUList):
         if attach:
             # Finally, append an HDU for this headerlet
             self.attach_to_file(fobj)
+            utils.updateNEXTENDKw(fobj)
         if close_dest:
             fobj.close()
 
@@ -2143,6 +2148,7 @@ class Headerlet(pyfits.HDUList):
 
         if attach:
             self.attach_to_file(fobj)
+            utils.updateNEXTENDKw(fobj)
 
         if close_dest:
             fobj.close()
@@ -2176,6 +2182,7 @@ class Headerlet(pyfits.HDUList):
             new_hlt = HeaderletHDU.fromheaderlet(self)
             new_hlt.header.update('extver', numhlt + 1)
             fobj.append(new_hlt)
+            utils.updateNEXTENDKw(fobj)
         else:
             message = "Headerlet %s cannot be attached to" % (self.hdrname)
             message += "observation %s" % (fname)
@@ -2310,7 +2317,7 @@ class Headerlet(pyfits.HDUList):
                     Distortion model in headerlet not the same as destination model
                     Headerlet model  : %s
                     Destination model: %s
-                    """ % (self[0].header['DISTNAME'], dname)
+                    """ % (self[0].header['DISTNAME'], dmodel)
                     logger.critical(message)
             return False
         else:
diff --git a/lib/stwcs/wcsutil/wcscorr.py b/lib/stwcs/wcsutil/wcscorr.py
index ac2484b..af04591 100644
--- a/lib/stwcs/wcsutil/wcscorr.py
+++ b/lib/stwcs/wcsutil/wcscorr.py
@@ -175,6 +175,8 @@ def init_wcscorr(input, force=False):
     # force an update now
     # set the verify flag to 'warn' so that it will always succeed, but still
     # tell the user if PyFITS detects any problems with the file as a whole
+    utils.updateNEXTENDKw(fimg)
+
     fimg.flush('warn')
 
     if need_to_close:
@@ -470,6 +472,8 @@ def restore_file_from_wcscorr(image, id='OPUS', wcskey=''):
                     pkey = key[:7]+wcskey
                 fimg[0].header.update(pkey,wcs_table.data.field(key)[erow])
 
+    utils.updateNEXTENDKw(fimg)
+
     # close the image now that the update has been completed.
     if close_image:
         fimg.close()