diff --git a/python/arrow/array.pxd b/python/arrow/array.pxd index 04dd8d182bcf6..482f8f796dd26 100644 --- a/python/arrow/array.pxd +++ b/python/arrow/array.pxd @@ -34,6 +34,7 @@ cdef class Array: DataType type cdef init(self, const shared_ptr[CArray]& sp_array) + cdef getitem(self, int i) cdef class BooleanArray(Array): diff --git a/python/arrow/array.pyx b/python/arrow/array.pyx index 8ebd01d1dbe73..b367e3b84a8b3 100644 --- a/python/arrow/array.pyx +++ b/python/arrow/array.pyx @@ -46,6 +46,17 @@ cdef class Array: def __get__(self): return self.sp_array.get().null_count() + def __iter__(self): + for i in range(len(self)): + yield self.getitem(i) + raise StopIteration + + def __repr__(self): + from arrow.formatting import array_format + type_format = object.__repr__(self) + values = array_format(self, window=10) + return '{0}\n{1}'.format(type_format, values) + def __len__(self): return self.sp_array.get().length() @@ -74,7 +85,10 @@ cdef class Array: while key < 0: key += len(self) - return scalar.box_arrow_scalar(self.type, self.sp_array, key) + return self.getitem(key) + + cdef getitem(self, int i): + return scalar.box_arrow_scalar(self.type, self.sp_array, i) def slice(self, start, end): pass diff --git a/python/arrow/scalar.pxd b/python/arrow/scalar.pxd index 15cdc956a2593..4e0a3647155a6 100644 --- a/python/arrow/scalar.pxd +++ b/python/arrow/scalar.pxd @@ -55,7 +55,7 @@ cdef class ListValue(ArrayValue): cdef: CListArray* ap - cdef _getitem(self, int i) + cdef getitem(self, int i) cdef class StringValue(ArrayValue): diff --git a/python/arrow/scalar.pyx b/python/arrow/scalar.pyx index 951ede2877690..72a280e334f4e 100644 --- a/python/arrow/scalar.pyx +++ b/python/arrow/scalar.pyx @@ -144,14 +144,19 @@ cdef class ListValue(ArrayValue): return self.ap.value_length(self.index) def __getitem__(self, i): - return self._getitem(i) + return self.getitem(i) + + def __iter__(self): + for i in range(len(self)): + yield self.getitem(i) + raise StopIteration cdef void _set_array(self, const shared_ptr[CArray]& sp_array): self.sp_array = sp_array self.ap = sp_array.get() self.value_type = box_data_type(self.ap.value_type()) - cdef _getitem(self, int i): + cdef getitem(self, int i): cdef int j = self.ap.offset(self.index) + i return box_arrow_scalar(self.value_type, self.ap.values(), j) @@ -161,7 +166,7 @@ cdef class ListValue(ArrayValue): list result = [] for j in range(len(self)): - result.append(self._getitem(j).as_py()) + result.append(self.getitem(j).as_py()) return result diff --git a/python/arrow/tests/test_array.py b/python/arrow/tests/test_array.py index 8eaa53352061b..ebd872c744e44 100644 --- a/python/arrow/tests/test_array.py +++ b/python/arrow/tests/test_array.py @@ -17,6 +17,7 @@ from arrow.compat import unittest import arrow +import arrow.formatting as fmt class TestArrayAPI(unittest.TestCase): @@ -24,3 +25,39 @@ class TestArrayAPI(unittest.TestCase): def test_getitem_NA(self): arr = arrow.from_pylist([1, None, 2]) assert arr[1] is arrow.NA + + def test_list_format(self): + arr = arrow.from_pylist([[1], None, [2, 3]]) + result = fmt.array_format(arr) + expected = """\ +[ + [1], + NA, + [2, + 3] +]""" + assert result == expected + + def test_string_format(self): + arr = arrow.from_pylist(['foo', None, 'bar']) + result = fmt.array_format(arr) + expected = """\ +[ + 'foo', + NA, + 'bar' +]""" + assert result == expected + + def test_long_array_format(self): + arr = arrow.from_pylist(range(100)) + result = fmt.array_format(arr, window=2) + expected = """\ +[ + 0, + 1, + ... + 98, + 99 +]""" + assert result == expected