Skip to content

Commit

Permalink
Basic array formatter, not tweaking too much for now
Browse files Browse the repository at this point in the history
  • Loading branch information
wesm committed Mar 8, 2016
1 parent df06ba1 commit 6777d11
Show file tree
Hide file tree
Showing 5 changed files with 62 additions and 5 deletions.
1 change: 1 addition & 0 deletions python/arrow/array.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
16 changes: 15 additions & 1 deletion python/arrow/array.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion python/arrow/scalar.pxd
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
11 changes: 8 additions & 3 deletions python/arrow/scalar.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -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 = <CListArray*> 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)

Expand All @@ -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

Expand Down
37 changes: 37 additions & 0 deletions python/arrow/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,47 @@

from arrow.compat import unittest
import arrow
import arrow.formatting as fmt


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

0 comments on commit 6777d11

Please sign in to comment.