Merge pull request #129 from GeminiDRSoftware/astrodata-slice

Allow to get an AstroData slice with 1 element

astrodata/core.py

@@ -71,7 +71,8 @@ class AstroData:
         'ut_date': 'DATE-OBS'
     }
 
-    def __init__(self, nddata=None, tables=None, phu=None, indices=None):
+    def __init__(self, nddata=None, tables=None, phu=None, indices=None,
+                 is_single=False):
         if nddata is None:
             nddata = []
         if not isinstance(nddata, (list, tuple)):
@@ -82,6 +83,10 @@ def __init__(self, nddata=None, tables=None, phu=None, indices=None):
         self._all_nddatas = nddata
         self._indices = indices
 
+        self.is_single = is_single
+        """ If this data provider represents a single slice out of a whole
+        dataset, return True. Otherwise, return False. """
+
         if tables is not None and not isinstance(tables, dict):
             raise ValueError('tables must be a dict')
         self._tables = tables or {}
@@ -289,14 +294,6 @@ def is_sliced(self):
         """
         return self._indices is not None
 
-    @property
-    def is_single(self):
-        """
-        If this data provider represents a single slice out of a whole dataset,
-        return True. Otherwise, return False.
-        """
-        return self._indices is not None and len(self._indices) == 1
-
     def is_settable(self, attr):
         """Return True if the attribute is meant to be modified."""
         if self.is_sliced and attr in {'path', 'filename'}:
@@ -482,8 +479,12 @@ def __getitem__(self, idx):
         if self._indices:
             indices = [self._indices[i] for i in indices]
 
-        obj = self.__class__(self._all_nddatas, tables=self._tables,
-                             phu=self.phu, indices=indices)
+        is_single = not isinstance(idx, (tuple, slice))
+        obj = self.__class__(self._all_nddatas,
+                             tables=self._tables,
+                             phu=self.phu,
+                             indices=indices,
+                             is_single=is_single)
         obj._path = self.path
         obj._orig_filename = self.orig_filename
         return obj

astrodata/tests/test_fits.py

@@ -211,6 +211,27 @@ def test_slice(GMOSN_SPECT):
         ext[1]
 
 
+@pytest.mark.dragons_remote_data
+def test_slice_single_element(GMOSN_SPECT):
+    ad = astrodata.open(GMOSN_SPECT)
+    assert ad.is_sliced is False
+
+    metadata = ('SCI', 2)
+
+    ext = ad[1:2]
+    assert ext.is_single is False
+    assert ext.is_sliced is True
+    assert ext.indices == [1]
+    assert isinstance(ext.data, list) and len(ext.data) == 1
+
+    ext = ext[0]
+    assert ext.id == 2
+    assert ext.is_single is True
+    assert ext.is_sliced is True
+    assert ext.hdr['EXTNAME'] == metadata[0]
+    assert ext.hdr['EXTVER'] == metadata[1]
+
+
 @pytest.mark.dragons_remote_data
 def test_slice_multiple(GMOSN_SPECT):
     ad = astrodata.open(GMOSN_SPECT)