diff --git a/sparsity/sparse_frame.py b/sparsity/sparse_frame.py index 9574e5b..6f015c6 100644 --- a/sparsity/sparse_frame.py +++ b/sparsity/sparse_frame.py @@ -789,7 +789,7 @@ def drop_duplicate_idx(self, **kwargs): Returns ------- - dropeed: SparseFrame + dropped: SparseFrame """ mask = ~self.index.duplicated(**kwargs) return SparseFrame(self.data[mask], index=self.index.values[mask], @@ -798,7 +798,10 @@ def drop_duplicate_idx(self, **kwargs): def __getitem__(self, item): if item is None: raise ValueError('Cannot label index with a null key.') - if not isinstance(item, (tuple, list)): + if not isinstance(item, (pd.Series, np.ndarray, pd.Index, list, + tuple)): + # TODO: tuple probably should be a separate case as in Pandas + # where it is used with Multiindex item = [item] if len(item) > 0: return self.reindex_axis(item, axis=1) diff --git a/sparsity/test/test_sparse_frame.py b/sparsity/test/test_sparse_frame.py index 3295c95..44b21f0 100644 --- a/sparsity/test/test_sparse_frame.py +++ b/sparsity/test/test_sparse_frame.py @@ -6,6 +6,7 @@ import numpy as np import pandas as pd +import pandas.testing as pdt import pytest from moto import mock_s3 from scipy import sparse @@ -575,6 +576,7 @@ def test_npz_io_s3(complex_example): def test_getitem(): id_ = np.identity(10) sf = SparseFrame(id_, columns=list('abcdefghij')) + assert sf['a'].data.todense()[0] == 1 assert sf['j'].data.todense()[9] == 1 assert np.all(sf[['a', 'b']].data.todense() == np.asmatrix(id_[:, [0, 1]])) @@ -588,6 +590,13 @@ def test_getitem(): with pytest.raises(ValueError): sf[None] + idx = pd.Index(list('abc')) + pdt.assert_index_equal(idx, sf[idx].columns) + pdt.assert_index_equal(idx, sf[idx.to_series()].columns) + pdt.assert_index_equal(idx, sf[idx.tolist()].columns) + pdt.assert_index_equal(idx, sf[tuple(idx)].columns) + pdt.assert_index_equal(idx, sf[idx.values].columns) + def test_vstack(): frames = []