BUG: bounds check in Cython get_value_at method, close #495

Author: wesm

pandas/core/series.py

@@ -260,6 +260,8 @@ def __getitem__(self, key):
                 except KeyError, e1:
                     try:
                         return _gin.get_value_at(self, key)
+                    except IndexError:
+                        raise
                     except Exception, _:
                         pass
                     raise e1

pandas/src/engines.pyx

@@ -6,33 +6,11 @@ cnp.import_array()
 
 cimport util
 
-cpdef inline object get_value_at(ndarray arr, object loc):
-    cdef:
-        Py_ssize_t i
-        void* data_ptr
-    if util.is_float_object(loc):
-        casted = int(loc)
-        if casted == loc:
-            loc = casted
-    i = <Py_ssize_t> loc
-    if i < 0:
-        i += cnp.PyArray_SIZE(arr)
-    data_ptr = cnp.PyArray_GETPTR1(arr, i)
-    return cnp.PyArray_GETITEM(arr, data_ptr)
-
-cpdef inline set_value_at(ndarray arr, object loc, object value):
-    cdef:
-        Py_ssize_t i
-    if util.is_float_object(loc):
-        casted = int(loc)
-        if casted == loc:
-            loc = casted
-    i = <Py_ssize_t> loc
-    if i < 0:
-        i += cnp.PyArray_SIZE(arr)
-
-    util.assign_value_1d(arr, i, value)
+def get_value_at(ndarray arr, object loc):
+    return util.get_value_at(arr, loc)
 
+def set_value_at(ndarray arr, object loc, object val):
+    return util.set_value_at(arr, loc, val)
 
 cdef class IndexEngine:
 
@@ -45,7 +23,7 @@ cdef class IndexEngine:
             void* data_ptr
 
         loc = self.get_loc(key)
-        return get_value_at(arr, loc)
+        return util.get_value_at(arr, loc)
 
     cpdef set_value(self, ndarray arr, object key, object value):
         '''
@@ -56,7 +34,7 @@ cdef class IndexEngine:
             void* data_ptr
 
         loc = self.get_loc(key)
-        set_value_at(arr, loc, value)
+        util.set_value_at(arr, loc, value)
 
 cdef class DictIndexEngine(IndexEngine):
     '''

pandas/src/util.pxd

@@ -8,29 +8,38 @@ cdef extern from "numpy_helper.h":
     inline int is_string_object(object)
     inline int assign_value_1d (ndarray, Py_ssize_t, object) except -1
 
-cpdef inline object get_value_at(ndarray arr, object loc):
+
+cdef inline object get_value_at(ndarray arr, object loc):
     cdef:
-        Py_ssize_t i
+        Py_ssize_t i, sz
         void* data_ptr
     if is_float_object(loc):
         casted = int(loc)
         if casted == loc:
             loc = casted
     i = <Py_ssize_t> loc
+    sz = cnp.PyArray_SIZE(arr)
+
     if i < 0:
-        i += cnp.PyArray_SIZE(arr)
+        i += sz
+    elif i >= sz:
+        raise IndexError('index out of bounds')
     data_ptr = cnp.PyArray_GETPTR1(arr, i)
     return cnp.PyArray_GETITEM(arr, data_ptr)
 
-cpdef inline set_value_at(ndarray arr, object loc, object value):
+cdef inline set_value_at(ndarray arr, object loc, object value):
     cdef:
-        Py_ssize_t i
+        Py_ssize_t i, sz
     if is_float_object(loc):
         casted = int(loc)
         if casted == loc:
             loc = casted
     i = <Py_ssize_t> loc
+    sz = cnp.PyArray_SIZE(arr)
+
     if i < 0:
-        i += cnp.PyArray_SIZE(arr)
+        i += sz
+    elif i >= sz:
+        raise IndexError('index out of bounds')
 
     assign_value_1d(arr, i, value)

pandas/tests/test_series.py

@@ -355,6 +355,10 @@ def test_getitem_setitem_boolean_corner(self):
         self.assertRaises(Exception, ts.ix.__getitem__, mask_shifted)
         self.assertRaises(Exception, ts.ix.__setitem__, mask_shifted, 1)
 
+    def test_getitem_out_of_bounds(self):
+        # don't segfault, GH #495
+        self.assertRaises(IndexError, self.ts.__getitem__, len(self.ts))
+
     def test_slice(self):
         numSlice = self.series[10:20]
         numSliceEnd = self.series[-10:]
@@ -531,7 +535,7 @@ def test_to_string(self):
         # pass float_format
         format = '%.4f'.__mod__
         result = self.ts.to_string(float_format=format)
-        result = [x.split()[1] for x in result.split('\n')[:-1]]
+        result = [x.split()[1] for x in result.split('\n')]
         expected = [format(x) for x in self.ts]
         self.assertEqual(result, expected)
 

setup.py

@@ -290,6 +290,7 @@ def srcpath(name=None, suffix='.pyx', subdir='src'):
 if suffix == '.pyx':
     tseries_depends = [srcpath(f, suffix='.pyx')
                        for f in tseries_depends]
+    tseries_depends.append('util.pxd')
 else:
     tseries_depends = []