Merge 446408c into d4a19ff

AUTHORS.rst

@@ -36,6 +36,7 @@ Alphabetical list of contributors
 * Brigitta Sipocz (@bsipocz)
 * Connor Stotts (@stottsco)
 * Ole Streicher (@olebole)
+* JVSN Reddy (@janga1997)
 * Erik Tollerud (@eteq)
 * Zè Vinícius (@mirca)
 * Josh Walawender (@joshwalawender)

CHANGES.rst

@@ -5,6 +5,7 @@
 New Features
 ^^^^^^^^^^^^
 
+- added feature to specify extension when calling ImageFileCollection [#463]
 
 Other Changes and Additions
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^

ccdproc/image_collection.py

@@ -58,14 +58,19 @@ class ImageFileCollection(object):
         The filenames are assumed to be in ``location``.
         Default is ``None``.
 
+    hdu: str or int, optional
+        Name/Index of the extension, which specifies the extension to extract data
+        from all files in a given directory.
+        Default is 0.
+
     Raises
     ------
     ValueError
         Raised if keywords are set to a combination of '*' and any other
         value.
     """
     def __init__(self, location=None, keywords=None, info_file=None,
-                 filenames=None):
+                 filenames=None, hdu=0):
         self._location = location
         self._filenames = filenames
         self._files = []
@@ -109,6 +114,8 @@ def __init__(self, location=None, keywords=None, info_file=None,
         # is always *some* value.
         self._all_keywords = False
 
+        self._hdu = hdu
+
         if keywords:
             self.keywords = keywords
 
@@ -363,7 +370,8 @@ def _add_val_to_dict(key, value, tbl_dict, n_previous, missing_marker):
             summary = input_summary
             n_previous = len(summary['file'])
 
-        h = fits.getheader(file_name)
+        h = fits.getheader(file_name, self._hdu)
+
         assert 'file' not in h
 
         # Try opening header before this so that file name is only added if
@@ -694,10 +702,13 @@ def _generator(self, return_type,
 
             file_name = path.basename(full_path)
 
+            hdu_index = hdulist.index_of(self._hdu)
+            ccd_kwargs.setdefault('hdu', hdu_index)
+
             return_options = {
-                    'header': lambda: hdulist[0].header,
-                    'hdu': lambda: hdulist[0],
-                    'data': lambda: hdulist[0].data,
+                    'header': lambda: hdulist[hdu_index].header,
+                    'hdu': lambda: hdulist[hdu_index],
+                    'data': lambda: hdulist[hdu_index].data,
                     'ccd': lambda: fits_ccddata_reader(full_path, **ccd_kwargs)
                     }
             try:

ccdproc/tests/test_image_collection.py

@@ -148,6 +148,26 @@ def test_hdus_masking(self, triage_setup):
         new_data = np.array(collection.summary_info)
         assert (new_data == old_data).all()
 
+    def test_multiple_extensions(self, triage_setup):
+        ext1 = fits.PrimaryHDU()
+        ext2 = fits.ImageHDU(name='MASK')
+        hdulist = fits.hdu.hdulist.HDUList([ext1, ext2])
+
+        hdulist.writeto(os.path.join(triage_setup.test_dir, 'multi-extension.fits'))
+        ic2 = image_collection.ImageFileCollection(triage_setup.test_dir, keywords='*',
+            filenames=['multi-extension.fits'], hdu='MASK')
+
+        ic1 = image_collection.ImageFileCollection(triage_setup.test_dir, keywords='*',
+            filenames=['multi-extension.fits'], hdu=0)
+
+        column2 = ic2.summary_info.colnames
+        column1 = ic1.summary_info.colnames
+
+        list1 = [key.lower() for key in ext2.header]
+        list2 = ic2.summary_info.colnames[1:]
+
+        assert list1 == list2
+
     def test_headers(self, triage_setup):
         collection = image_collection.ImageFileCollection(location=triage_setup.test_dir,
                                              keywords=['imagetyp'])