TST: indexing testing with minor Series.__getitem__ refactoring

Author: wesm

pandas/core/format.py

@@ -101,10 +101,8 @@ def _get_formatted_values(self):
                     fmt_values.append(' %s' % self.formatter(v))
                 elif is_float[i]:
                     fmt_values.append(float_format(v))
-                elif not leading_space:
-                    fmt_values.append(' %s' % self.formatter(v))
                 else:
-                    fmt_values.append(self.formatter(v))
+                    fmt_values.append(' %s' % self.formatter(v))
         else:
             fmt_values = _format_fixed_width(self.series.values,
                                              self.formatter)

pandas/core/index.py

@@ -10,7 +10,7 @@
 from pandas.util.decorators import cache_readonly
 import pandas.core.common as com
 import pandas._tseries as lib
-import pandas._engines as _engines
+import pandas._engines as _gin
 
 __all__ = ['Index']
 
@@ -24,6 +24,8 @@ def wrapper(self, other):
         return func(other)
     return wrapper
 
+class InvalidIndexError(Exception):
+    pass
 
 class Index(np.ndarray):
     """
@@ -146,8 +148,8 @@ def _cleanup(self):
     def _engine(self):
         import weakref
         # property, for now, slow to look up
-        return _engines.DictIndexEngine(weakref.ref(self),
-                                        self._map_indices)
+        return _gin.DictIndexEngine(weakref.ref(self),
+                                    self._map_indices)
 
     def _get_level_number(self, level):
         if not isinstance(level, int):
@@ -469,12 +471,31 @@ def get_loc(self, key):
         """
         return self._engine.get_loc(key)
 
-    def get_value(self, arr, key):
+    def get_value(self, series, key):
         """
         Fast lookup of value from 1-dimensional ndarray. Only use this if you
         know what you're doing
         """
-        return self._engine.get_value(arr, key)
+        try:
+            return self._engine.get_value(series, key)
+        except KeyError, e1:
+            if self.inferred_type == 'integer':
+                raise
+
+            try:
+                return _gin.get_value_at(series, key)
+            except IndexError:
+                raise
+            except TypeError:
+                # generator/iterator-like
+                if hasattr(key, 'next'):
+                    raise InvalidIndexError(key)
+                else:
+                    raise e1
+            except Exception:  # pragma: no cover
+                raise e1
+        except TypeError:
+            raise InvalidIndexError(key)
 
     def set_value(self, arr, key, value):
         """
@@ -1091,6 +1112,40 @@ def has_duplicates(self):
 
         return False
 
+    def get_value(self, series, key):
+        # somewhat broken encapsulation
+        from pandas.core.indexing import _maybe_droplevels
+        from pandas.core.series import Series
+
+        # Label-based
+        try:
+            return self._engine.get_value(series, key)
+        except KeyError, e1:
+            try:
+                # TODO: what if a level contains tuples??
+                loc = self.get_loc(key)
+                new_values = series.values[loc]
+                new_index = self[loc]
+                new_index = _maybe_droplevels(new_index, key)
+                return Series(new_values, index=new_index, name=series.name)
+            except KeyError:
+                pass
+
+            try:
+                return _gin.get_value_at(series, key)
+            except IndexError:
+                raise
+            except TypeError:
+                # generator/iterator-like
+                if hasattr(key, 'next'):
+                    raise InvalidIndexError(key)
+                else:
+                    raise e1
+            except Exception:  # pragma: no cover
+                raise e1
+        except TypeError:
+            raise InvalidIndexError(key)
+
     def get_level_values(self, level):
         """
         Return vector of label values for requested level, equal to the length
@@ -1641,22 +1696,31 @@ def get_loc_level(self, key, level=0):
             else:
                 indexer = None
                 for i, k in enumerate(key):
-                    if k is None:
-                        continue
+                    if not isinstance(k, slice):
+                        k = self._get_level_indexer(k, level=i)
+                        if isinstance(k, slice):
+                            # everything
+                            if k.start == 0 and k.stop == len(self):
+                                k = slice(None, None)
+                        else:
+                            k_index = k
 
                     if isinstance(k, slice):
                         if k == slice(None, None):
-                           continue
+                            continue
                         else:
-                            k_index = np.empty(len(self), dtype=bool)
-                            k_index[k] = True
-                    else:
-                        k_index = self._get_level_indexer(k, level=i)
+                            raise NotImplementedError
+                            # if self.levels[i].inferred_type == 'integer':
+                            #     raise NotImplementedError
+                            # k_index = np.zeros(len(self), dtype=bool)
+                            # k_index[k] = True
 
                     if indexer is None:
                         indexer = k_index
-                    else:
+                    else:  # pragma: no cover
                         indexer &= k_index
+                if indexer is None:
+                    indexer = slice(None, None)
                 return indexer
         else:
             return self._get_level_indexer(key, level=level)

pandas/core/series.py

@@ -8,6 +8,7 @@
 from itertools import izip
 import csv
 import operator
+import types
 
 from numpy import nan, ndarray
 import numpy as np
@@ -18,17 +19,16 @@
                                 _asarray_tuplesafe)
 from pandas.core.daterange import DateRange
 from pandas.core.format import SeriesFormatter
-from pandas.core.index import Index, MultiIndex, _ensure_index
-from pandas.core.indexing import _SeriesIndexer, _maybe_droplevels
+from pandas.core.index import (Index, MultiIndex, InvalidIndexError,
+                               _ensure_index)
+from pandas.core.indexing import _SeriesIndexer
 from pandas.util import py3compat
 from pandas.util.terminal import get_terminal_size
 import pandas.core.common as com
 import pandas.core.datetools as datetools
 import pandas.core.generic as generic
 import pandas.core.nanops as nanops
 import pandas._tseries as lib
-import pandas._engines as _gin
-
 from pandas.util.decorators import Appender, Substitution
 
 __all__ = ['Series', 'TimeSeries']
@@ -261,39 +261,17 @@ def ix(self):
         return self._ix
 
     def __getitem__(self, key):
-        index = self.index
-
-        # Label-based
         try:
-            return index._engine.get_value(self, key)
-        except KeyError, e1:
-            if isinstance(index, MultiIndex):
-                values = self.values
-                try:
-                    loc = index.get_loc(key)
-                    # TODO: what if a level contains tuples??
-                    new_index = index[loc]
-                    new_index = _maybe_droplevels(new_index, key)
-                    return Series(values[loc], index=new_index,
-                                  name=self.name)
-                except KeyError:
-                    pass
-
-            if index.inferred_type == 'integer':
-                raise
-
-            try:
-                return _gin.get_value_at(self, key)
-            except IndexError:
-                raise
-            except Exception, _:
-                pass
-            raise e1
-        except TypeError:
+            return self.index.get_value(self, key)
+        except InvalidIndexError:
             pass
+        except Exception:
+            raise
 
-        # boolean
+        if hasattr(key, 'next'):
+            key = list(key)
 
+        # boolean
         # special handling of boolean data with NAs stored in object
         # arrays. Since we can't represent NA with dtype=bool
         if _is_bool_indexer(key):
@@ -311,7 +289,7 @@ def _get_with(self, key):
             if isinstance(key, tuple):
                 return self._get_values_tuple(key)
 
-            if not isinstance(key, (list, np.ndarray)):
+            if not isinstance(key, (list, np.ndarray)):  # pragma: no cover
                 key = list(key)
 
             key_type = lib.infer_dtype(key)
@@ -348,7 +326,7 @@ def _get_values_tuple(self, key):
         # kludgearound
         new_index = result.index
         for i, k in reversed(list(enumerate(key))):
-            if k != slice(None, None):
+            if not isinstance(k, slice):
                 new_index = new_index.droplevel(i)
         result.index = new_index
 

pandas/tests/test_frame.py

@@ -2322,6 +2322,8 @@ def test_append_series_dict(self):
 
         series = df.ix[4]
         self.assertRaises(Exception, df.append, series)
+        series.name = None
+        self.assertRaises(Exception, df.append, series)
 
         result = df.append(series[::-1], ignore_index=True)
         expected = df.append(DataFrame({0 : series[::-1]},

pandas/tests/test_index.py

@@ -803,6 +803,31 @@ def test_get_loc_duplicates(self):
         index = Index([2, 2, 2, 2])
         self.assertRaises(Exception, index.get_loc, 2)
 
+    def test_get_loc_level(self):
+        index = MultiIndex(levels=[Index(range(4)),
+                                   Index(range(4)),
+                                   Index(range(4))],
+                           labels=[np.array([0, 0, 1, 2, 2, 2, 3, 3]),
+                                   np.array([0, 1, 0, 0, 0, 1, 0, 1]),
+                                   np.array([1, 0, 1, 1, 0, 0, 1, 0])])
+
+        loc = index.get_loc_level((0, 1))
+        expected = slice(1, 2)
+        self.assertEqual(loc, expected)
+
+        loc = index.get_loc_level((0, 1, 0))
+        expected = 1
+        self.assertEqual(loc, expected)
+
+        self.assertRaises(KeyError, index.get_loc_level, (2, 2))
+
+        index = MultiIndex(levels=[[2000], range(4)],
+                           labels=[np.array([0, 0, 0, 0]),
+                                   np.array([0, 1, 2, 3])])
+        result = index.get_loc_level((2000, slice(None, None)))
+        expected = slice(None, None)
+        self.assertEqual(result, expected)
+
     def test_slice_locs(self):
         df = tm.makeTimeDataFrame()
         stacked = df.stack()

pandas/tests/test_multilevel.py

@@ -175,6 +175,18 @@ def test_series_getitem(self):
         # key error
         self.assertRaises(KeyError, s.__getitem__, (2000, 3, 4))
 
+    def test_series_getitem_corner(self):
+        s = self.ymd['A']
+
+        # don't segfault, GH #495
+        # out of bounds access
+        self.assertRaises(IndexError, s.__getitem__, len(self.ymd))
+
+        # generator
+        result = s[(x > 0 for x in s)]
+        expected = s[s > 0]
+        assert_series_equal(result, expected)
+
     def test_series_setitem(self):
         s = self.ymd['A']
 
@@ -216,6 +228,24 @@ def test_xs_level_series(self):
         expected = self.frame.xs('two', level=1)['A']
         assert_series_equal(result, expected)
 
+        s = self.ymd['A']
+        result = s[2000, 5]
+        expected = self.ymd.ix[2000, 5]['A']
+        assert_series_equal(result, expected)
+
+        # not implementing this for now
+
+        self.assertRaises(NotImplementedError, s.__getitem__,
+                          (2000, slice(3, 4)))
+
+        # result = s[2000, 3:4]
+        # lv =s.index.get_level_values(1)
+        # expected = s[(lv == 3) | (lv == 4)]
+        # expected.index = expected.index.droplevel(0)
+        # assert_series_equal(result, expected)
+
+        # can do this though
+
     def test_fancy_2d(self):
         result = self.frame.ix['foo', 'B']
         expected = self.frame.xs('foo')['B']

pandas/tests/test_series.py

@@ -361,6 +361,14 @@ def test_getitem_boolean(self):
         assert_series_equal(result, expected)
         self.assert_(np.array_equal(result.index, s.index[mask]))
 
+    def test_getitem_generator(self):
+        gen = (x > 0 for x in self.series)
+        result = self.series[gen]
+        result2 = self.series[iter(self.series > 0)]
+        expected = self.series[self.series > 0]
+        assert_series_equal(result, expected)
+        assert_series_equal(result2, expected)
+
     def test_getitem_boolean_object(self):
         # using column from DataFrame
         s = self.series
@@ -737,6 +745,14 @@ def test_to_string_mixed(self):
                     '3    baz')
         self.assertEqual(result, expected)
 
+        s = Series(['foo', 5, 'bar', 'baz'])
+        result = s.to_string()
+        expected = ('0    foo\n'
+                    '1    5\n'
+                    '2    bar\n'
+                    '3    baz')
+        self.assertEqual(result, expected)
+
     def test_to_string_float_na_spacing(self):
         s = Series([0., 1.5678, 2., -3., 4.])
         s[::2] = np.nan