Merge pull request #238 from mommermi/data_slice_fix

better __getitem__ handling in data.DataClass

CHANGES

@@ -1,6 +1,9 @@
-0.3dev
+0.2.2dev
 ======
-
+- sbpy.data.DataClass.__getitem__ now always returns a new object of the same
+  class, unless a single field name is provided in which case an
+  astropy.Table.Column (no units provided) or astropy.units.Quantity
+  (units provided) is returned.
 
 This changelog tracks changes to sbpy starting with version v0.2.
 

docs/sbpy/data/dataclass.rst

@@ -300,7 +300,8 @@ Each of these columns can be accessed easily, for instance:
     [10.223423 10.233453 10.243452] deg
 
 which will return an `~astropy.units.quantity.Quantity` object if that
-column has a `~astropy.units.Unit` attached to it.
+column has a `~astropy.units.Unit` attached to it or a `~astropy.table.Column`
+otherwise.
 
 Similarly, if you are interested in the first set of observations in
 ``obs``, you can use:
@@ -311,7 +312,8 @@ Similarly, if you are interested in the first set of observations in
     --------- --------- ------------
     10.223423 -12.42123 2451523.6234
 
-which returns you a table with only the requested subset of the
+which returns you a new instance of the same class as your original
+objet with only the requested subset of the
 data. In order to retrieve RA from the second observation, you can
 combine both examples and do:
 
@@ -332,6 +334,13 @@ for instance:
     10.233453 -12.41562
     10.243452 -12.40435
 
+    >>> obs[:2] # doctest: +SKIP
+        ra       dec         t
+       deg       deg         d
+    --------- --------- ------------
+    10.223423 -12.42123 2451523.6234
+    10.233453 -12.41562 2451523.7345
+
     >>> obs[obs['ra'] <= 10.233453*u.deg] # doctest: +SKIP
         ra       dec         t
        deg       deg         d

sbpy/data/core.py

@@ -614,32 +614,31 @@ def __repr__(self):
 
     def __getitem__(self, ident):
         """Return columns or rows from data table(``self._table``); checks
-        for and may use alternative field names."""
-
-        # iterable
-        if isinstance(ident, (list, tuple, ndarray)):
-            if all([isinstance(i, str) for i in ident]):
-                # list of column names
+        for and may use alternative field names. This method will always return
+        an instance of __class__, except in the case when a field name is
+        requested (then return an `astropy.table.Column` if no units are
+        provided or a `astropy.units.Quantity` if units are provided)."""
+
+        # slices, iterables consisting of booleans and integers, and integer
+        # indices are all treated in the same way and are required to return
+        # a new __class__ object; only have to treat string identifiers
+        # separately in that those have to be checked for conversions
+        # and translations
+
+        # list of field names
+        if (isinstance(ident, (list, tuple, ndarray)) and
+            all([isinstance(i, str) for i in ident])):
                 self = self._convert_columns(ident)
                 newkeylist = [self._translate_columns(i)[0] for i in ident]
                 ident = newkeylist
-                # return as new DataClass object
-                return self.from_table(self._table[ident])
-            # ignore lists of boolean (masks)
-            elif all([isinstance(i, bool) for i in ident]):
-                pass
-            # ignore lists of integers
-            elif all([isinstance(i, int) for i in ident]):
-                pass
-        # individual strings
+        # individual field names
         elif isinstance(ident, str):
             self = self._convert_columns(ident)
             ident = self._translate_columns(ident)[0]
-        elif isinstance(ident, int):
-            return self.from_table(self._table[ident])
+            return self._table[ident]
 
-        # return as element from self_table
-        return self._table[ident]
+        # return as new instance of this class for all other identifiers
+        return self.from_table(self._table[ident])
 
     def __setitem__(self, *args):
         """Refer cls.__setitem__ to self._table"""

sbpy/data/tests/test_dataclass.py

@@ -5,7 +5,7 @@
 from copy import deepcopy
 from numpy import array
 import astropy.units as u
-from astropy.table import QTable
+from astropy.table import QTable, Column
 from ..core import DataClass, conf, DataClassError
 
 
@@ -257,20 +257,41 @@ def test_get_set():
 
     data = DataClass.from_dict(
         OrderedDict((('a', [1, 2, 3]),
-                     ('b', [4, 5, 6]),
+                     ('b', [4, 5, 6]*u.m),
                      ('c', [7, 8, 8]))))
 
     # get a single column
-    assert len(data['a']) == 3
+    x = data['a']
+    assert len(x) == 3
+    assert isinstance(x, Column)
+    x = data['b']
+    assert len(x) == 3
+    assert isinstance(x, u.Quantity)
+
+    # get a list of columns
+    x = data[['a', 'c']]
+    assert len(x.field_names) == 2
+    assert isinstance(x, DataClass)
 
     # mask rows
     masked = data[[True, False, False]]
     assert len(masked) == 1
-    assert masked['b'][0] == 4
+    assert masked['b'][0] == 4*u.m
+    assert isinstance(masked, DataClass)
+
+    # get single row
+    shortened = data[1]
+    assert shortened['a'] == 2
+    assert isinstance(shortened, DataClass)
 
     # get list of rows
     shortened = data[[0, 1]]
     assert len(shortened) == 2
+    assert isinstance(shortened, DataClass)
+
+    # get slice
+    assert len(shortened) == 2
+    assert isinstance(shortened, DataClass)
 
     # modify an existing column
     data['a'][:] = [0, 0, 0]
@@ -282,10 +303,12 @@ def test_get_set():
     # add non-existing column using set
     data['z'] = 3
     assert len(data['z'] == 3)
+    assert isinstance(data, DataClass)
 
     # modify existing column using set
     data['z'] = 2
     assert data['z'][1] == 2
+    assert isinstance(data, DataClass)
 
 
 def test_units():