Allow group specifications as integers or tuples (#1612)

CHANGELOG.rst

@@ -25,6 +25,11 @@ number of the code change for that issue.  These PRs can be viewed at:
 
 - Fixed projection cell identification in overlapping regions. [#1572]
 
+- Fixed a bug in processing of the ``group`` argument due to which the code
+  would crash when ``group`` would be an integer number or a list of numbers.
+  Also, added support for specifying extensions as tuples of
+  ``(extname, extver)``. [#1612]
+
 - Force the version of matplotlib to be <= 3.6.3 as the newer versions of
   the library cause problems with the calcloud preview generation.  This
   is a temporary restriction.

drizzlepac/imageObject.py

@@ -1038,26 +1038,11 @@ def __init__(self,filename, output=None, group=None,inmemory=False):
         fimg.close()
         del fimg
 
-        if group not in [None,'']:
-            # Only use selected chip
-            if ',' in group:
-                group_id = group.split(',')
-                if group_id[0].isalpha(): # user specified a specific extname,extver
-                    self.group = [int(group_id[1])]
-                else: # user specified a list of extension numbers to process
-                    self.group = []
-                    for grp in group_id:
-                        # find extname/extver which corresponds to this extension number
-                        group_extname = self._image[int(grp)].header['EXTNAME']
-                        group_extver = self._image[int(grp)].header['EXTVER']
-                        self.group.append(group_extver)
-            else:
-                # find extname/extver which corresponds to this extension number
-                group_extver = self._image[int(group)].header['EXTVER']
-                self.group = [int(group_extver)]
-        else:
-            # Use all chips
-            self.group = None
+        self.group = util._parse_ext_spec(
+            hdulist=self._image,
+            extno=group,
+            extname=self.scienceExt
+        )
 
         if not self._isSimpleFits:
 

drizzlepac/processInput.py

@@ -375,22 +375,11 @@ def applyContextPar(imageObjectList,contextpar):
 def _getInputImage (input, output=None, group=None):
     """ Factory function to return appropriate imageObject class instance"""
     # extract primary header and SCI,1 header from input image
-    sci_ext = 'SCI'
-    if group in [None,'']:
-        exten = '[sci,1]'
-        phdu = fits.getheader(input, memmap=False)
-    else:
-        # change to use fits more directly here?
-        if group.find(',') > 0:
-            grp = group.split(',')
-            if grp[0].isalpha():
-                grp = (grp[0],int(grp[1]))
-            else:
-                grp = int(grp[0])
-        else:
-            grp = int(group)
-        phdu = fits.getheader(input, memmap=False)
-        phdu.extend(fits.getheader(input, ext=grp, memmap=False))
+    phdu = fits.getheader(input, memmap=False)
+    grp = util._parse_ext_spec(input, extno=group, extname=None)
+    exten = '[sci,1]'
+    if grp is not None:
+        phdu.extend(fits.getheader(input, ext=grp[0], memmap=False))
 
     # Extract the instrument name for the data that is being processed by Multidrizzle
     _instrument = phdu['INSTRUME']

drizzlepac/util.py

@@ -1263,3 +1263,97 @@ def help(file=None):
     module['help'] = help
 
     return getHelpAsString(docstring=True, show_ver=False)
+
+
+def _parse_ext_spec(hdulist, extno, extname=None):
+    # process FITS extension number(s) and return a list of extension versions
+    # for the input extname.
+    if extno in [None, '']:
+        return None
+
+    groups = None
+
+    if isinstance(hdulist, str):
+        hdulist = fits.open(hdulist)
+        close_hdulist = True
+    else:
+        close_hdulist = False
+
+    def close_image_file():
+        if close_hdulist:
+            hdulist.close()
+
+    if isinstance(extno, str):
+        exts = list(map(str.strip, extno.split(',')))
+        if len(exts) == 2 and exts[0].isalpha():
+            # user specified a specific extname, extver
+            groups = [(exts[0], int(exts[1]))]
+        else:
+            # user specified a list of extension numbers:
+            extno = exts
+
+    elif isinstance(extno, int):
+        # user provided a specific extname, extver in the form of a tuple
+        extno = [extno]
+
+    elif isinstance(extno, tuple):
+        # user provided a specific extname, extver in the form of a tuple
+        groups = [extno]
+
+    if groups is None and isinstance(extno, list):
+        # user specified a list of extension numbers if the form of a list
+        # of integers or strings convertable to integers:
+        groups = list(map(int, extno))
+        if max(groups) >= len(hdulist):
+            close_image_file()
+            raise ValueError(
+                "Requested FITS extension number is larger than the "
+                "number extensions in the input file."
+            )
+        try:
+            groups = [
+                (
+                    hdulist[g].header['extname'].strip().lower(),
+                    hdulist[g].header['extver']
+                )
+                for g in groups
+            ]
+        except KeyError:
+            close_image_file()
+            raise ValueError(
+                f"The 'group' parameter must correspond to '{extname}' "
+                "FITS extensions of HST data files."
+            )
+
+    elif groups is None:
+        close_image_file()
+        raise TypeError("Unsupported group format.")
+
+    if extname is not None:
+        extname = extname.strip().lower()
+
+        # make sure actual extension name is the same as the desired extname
+        # and that the extension exists:
+        for extn, extv in groups:
+            if extn != extname:
+                close_image_file()
+                raise ValueError(
+                    "Groups must correspond to FITS extensions with "
+                    f"'{extname}' extension name."
+                )
+            if (extn, extv) not in hdulist:
+                if hdulist.filename is None:
+                    f = "in-memory file."
+                else:
+                    f = f"file '{hdulist.filename()}'."
+                close_image_file()
+                raise ValueError(
+                    f"Requested group ({extn}, {extv}) does not exist in the "
+                    f"input {f}"
+                )
+
+    close_image_file()
+
+    # return a list of extension versions corresponding to input 'extname'
+    extvers = [extv for _, extv in groups]
+    return extvers