diff --git a/dios/dios.py b/dios/dios.py
index 44aaad4423e4af68e00a37549f7be8c906edc983..35efd8a05b7a29d8b1bd0135cb673c322c2c6854 100644
--- a/dios/dios.py
+++ b/dios/dios.py
@@ -599,15 +599,33 @@ class _LocIndexer(_Indexer):
self._set_item = _dios._set_item
def __getitem__(self, key):
- rkey, cols = self._unpack_key(key)
- new = self._dios.copy_empty()
- for i, _ in enumerate(cols):
- c, r = cols[i], rkey[i]
- new[c] = self._data[c].loc[r]
- return new
+ rkey, cols, lowdim = self._unpack_key(key)
+ if is_scalar(rkey[0]):
+ return self._series(rkey, cols, lowdim)
+ elif lowdim:
+ return self._scalar(rkey[0], cols[0])
+ else:
+ new = self._dios.copy_empty()
+ for i, _ in enumerate(cols):
+ c, r = cols[i], rkey[i]
+ new[c] = self._data[c].loc[r]
+ return new
+
+ def _series(self, rkey, cols, lowdim):
+ if lowdim:
+ return self._scalar(rkey[0], cols[0])
+ new = pd.Series()
+ for c in cols:
+ try:
+ new[c] = self._data[c].loc[rkey]
+ except KeyError:
+ new[c] = np.nan
+
+ def _scalar(self, r, c):
+ return self._data[c].loc[r]
def __setitem__(self, key, value):
- ixs, keys = self._unpack_key(key)
+ ixs, keys, _ = self._unpack_key(key)
gen = self._unpack_value(keys, ixs, value)
for tup in gen:
self._set_item(*tup)
@@ -615,6 +633,7 @@ class _LocIndexer(_Indexer):
def _unpack_key(self, key):
# if we have a tuple, we have a rows- and a column-indexer
# if not, we only have a row-indexer and work on all columns
+ lowdim = False
if isinstance(key, tuple):
rkey, ckey, *fail = key
if fail:
@@ -628,21 +647,21 @@ class _LocIndexer(_Indexer):
raise ValueError("Cannot index with multidimensional key")
if isinstance(ckey, str):
cols = [ckey]
+ lowdim = True
elif isinstance(ckey, slice):
cols = self._col_slice_to_col_list(ckey)
else:
try:
- # list and bool list like
+ # list, boolean-list or series
cols, *_ = self._dios._unpack_key(ckey)
-
- except Exception:
- raise
+ except Exception as e:
+ raise e
else:
cols = list(self._data.keys())
rkey = key
# blowup
rkey = [rkey] * len(cols)
- return rkey, cols
+ return rkey, cols, lowdim
def _col_slice_to_col_list(self, cslice):
""" see here:
@@ -662,15 +681,33 @@ class _LocIndexer(_Indexer):
class _iLocIndexer(_Indexer):
def __getitem__(self, key):
- rkey, cols = self._unpack_key(key)
- new = self._dios.copy_empty()
- for i, _ in enumerate(cols):
- c, r = cols[i], rkey[i]
- new[c] = self._data[c].iloc[r]
- return new
+ rkey, cols, lowdim = self._unpack_key(key)
+ if is_scalar(rkey[0]):
+ return self._series(rkey, cols, lowdim)
+ elif lowdim:
+ return self._scalar(rkey[0], cols[0])
+ else:
+ new = self._dios.copy_empty()
+ for i, _ in enumerate(cols):
+ c, r = cols[i], rkey[i]
+ new[c] = self._data[c].iloc[r]
+ return new
+
+ def _series(self, rkey, cols, lowdim):
+ if lowdim:
+ return self._scalar(rkey[0], cols[0])
+ new = pd.Series()
+ for c in cols:
+ try:
+ new[c] = self._data[c].iloc[rkey]
+ except KeyError:
+ new[c] = np.nan
+
+ def _scalar(self, r, c):
+ return self._data[c].iloc[r]
def __setitem__(self, key, value):
- ixs, keys = self._unpack_key(key)
+ ixs, keys, _ = self._unpack_key(key)
gen = self._unpack_value(keys, ixs, value)
for tup in gen:
self._set_item_positional(*tup)
@@ -689,6 +726,7 @@ class _iLocIndexer(_Indexer):
def _unpack_key(self, key):
# if we have a tuple, we have a rows- and a column-indexer
# if not, we only have a row-indexer and work on all columns
+ lowdim = False
if isinstance(key, tuple):
rkey, ckey, *fail = key
if fail:
@@ -701,6 +739,7 @@ class _iLocIndexer(_Indexer):
if is_integer(ckey):
self._check_keys([ckey])
cols = self._integers_to_col_list([ckey])
+ lowdim = True
elif isinstance(ckey, slice):
cols = self._col_slice_to_col_list(ckey)
elif is_list_like(ckey) and not is_nested_list_like(ckey):
@@ -721,7 +760,7 @@ class _iLocIndexer(_Indexer):
# blowup
rkey = [rkey] * len(cols)
- return rkey, cols
+ return rkey, cols, lowdim
def _check_keys(self, keys):
bound = len(self._data)
diff --git a/test/test__getsetitem__.py b/test/test__getsetitem__.py
index bc0019794b36a52e73154c103fd758fffc939c9a..56a2d7ff0a30d965c2da68bb0cdfb695d4afc51b 100644
--- a/test/test__getsetitem__.py
+++ b/test/test__getsetitem__.py
@@ -12,27 +12,34 @@ d1 = DictOfSeries(data=dict(a=s1.copy(), b=s2.copy(), c=s3.copy(), d=s4.copy()))
@pytest.mark.parametrize(('idxer', 'exp'), [('a', s1), ('c', s3)])
def test__getitem_single(idxer, exp):
a = d1[idxer]
+ b = d1.loc[:, idxer]
assert isinstance(a, pd.Series)
+ assert isinstance(b, pd.Series)
assert (a == exp).all()
+ assert (b == exp).all()
-@pytest.mark.parametrize(('idxer', 'exp'), [('a', s1), ('c', s3)])
-def test__getitem_single_loc(idxer, exp):
- a = d1.loc[:, idxer]
- assert isinstance(a, DictOfSeries)
- exp = DictOfSeries(data={idxer: exp})
- assert (a == exp).all().all()
+@pytest.mark.parametrize(('idxer', 'exp'), [((1, 'a'), s1), ((3, 'c'), s3)])
+def test__getitem_scalar_loc(idxer, exp):
+ a = d1.loc[idxer]
+ assert is_scalar(a)
+ assert a == exp.loc[idxer[0]]
@pytest.mark.parametrize(('idxer', 'exp'), [(0, s1), (1, s2), (2, s3), (3, s4),
(-1, s4), (-2, s3), (-3, s2), (-4, s1)])
def test__getitem_single_iloc(idxer, exp):
a = d1.iloc[:, idxer]
- assert isinstance(a, DictOfSeries)
- exp = DictOfSeries(data={exp.name: exp})
- assert (a == exp).all().all()
+ assert isinstance(a, pd.Series)
+ assert (a == exp).all()
+@pytest.mark.parametrize(('idxer', 'exp'), [((1,0), s1), ((3,-2), s3), ((-1,-1), s4)])
+def test__getitem_scalar_iloc(idxer, exp):
+ a = d1.iloc[idxer]
+ assert is_scalar(a)
+ assert a == exp.iloc[idxer[0]]
+
@pytest.mark.parametrize('idxer', ['x', '1', 1, None, ])
def test__getitem_single_fail(idxer):
with pytest.raises(KeyError):
@@ -51,7 +58,6 @@ def test__getitem_single_iloc_fail(idxer):
a = d1.iloc[:, idxer]
-
BLIST = [True, False, False, True]
BOOLIDXER = [BLIST, pd.Series(BLIST), d1.copy() > 10]
INDEXERS = [['a'], ['a', 'c'], pd.Series(['a', 'c'])] + BOOLIDXER