def _sanitize_rows(
rows: Iterable,
indices: Union[int, Mapping] = None,
result_index: "Index" = None,
) -> np.ndarray:
"""Sanitize rows passed to slice"""
from ..base import c
if is_integer(indices):
rows = Collection(*rows, pool=indices)
if rows.error:
raise rows.error from None
return np.array(rows, dtype=int)
out = []
if any(isinstance(row, SeriesGroupBy) for row in rows):
rows = c(*rows)
for key in result_index:
idx = dict_get(indices, key)
if idx.size == 0:
continue
gidx = dict_get(rows.grouper.indices, key)
out.extend(idx.take(rows.obj.take(gidx)))
else:
for key in result_index:
idx = dict_get(indices, key)
if idx.size == 0:
continue
grows = Collection(*rows, pool=idx.size)
if grows.error:
raise grows.error from None
out.extend(idx.take(grows))
return np.array(out, dtype=int)
def test_negate_collection():
# cs = -Collection(1, 2)
# assert cs == [-1, -2]
cs = -Collection(0, 1, pool=2)
# pool: 0, 1 with 1, 2, we get 1
# -2, -1
# 0, 1
assert list(cs) == [0, 1]
cs = -Collection(1, 2, pool=[])
assert cs == []
assert cs.unmatched == {1, 2}
assert "Negated" in repr(cs)
cs = -Collection(list("abc"), pool=list("bcd"))
# 'abc' mapping to 'bcd' => [0,1]
# Negated([0,1], pool=3) => [-2,-1] => [1,2]
assert list(cs) == [1, 2]
assert cs.unmatched == {"a"}
# cannot negate strings
with pytest.raises(TypeError):
-Collection("a")
cs = -Collection("c", pool=["a", "b"])
assert cs == []
assert cs.unmatched == {"c"}
def test_mixed_collections(s):
cs = Collection(1, 2, s[3:5], -Collection(1, 2))
assert cs == [1, 2, 3, 4, -1, -2]
cs = Collection(
3, Collection(5), s[6:7], -Collection(1), pool=10
)
assert list(cs) == [3, 5, 6, 9]
def test_mixed_negatives_and_positives():
cs = Collection(-Collection(1, 3), 4)
assert cs == [-1, -3, 4]
# 0 1 2 3 4
# -3 -1
# 3
cs = Collection(-Collection(1, 3), 3, pool=5)
assert list(cs) == [4, 2, 3]
def test_slice_with_int_pool(s):
slc = Collection(s[1:3], pool=4)
assert list(slc) == [1, 2]
slc = Collection(s[:3], pool=4)
assert list(slc) == [0, 1, 2]
slc = Collection(s[1:], pool=4)
assert list(slc) == [1, 2, 3]
slc = Collection(s[:], pool=4)
assert list(slc) == [0, 1, 2, 3]
def test_invert_collection():
cs = ~Collection(1, 2)
# cannot expand without pool
assert cs == []
assert "Inverted" in repr(cs)
# non-index gets excluded literally, but returns indexes
cs = ~Collection("a", "b", pool=list("abc"))
assert cs == [2]
cs = Inverted(0, 1, 2, 3, pool=5)
assert cs == [4]
def test_collections_flattened():
cs = Collection(Collection(3, 4))
assert cs == [3, 4]
cs = Collection(1, 2, Collection(3, 4))
assert cs == [1, 2, 3, 4]
cs = Collection(1, 2, Collection(3, Collection(4)))
assert cs == [1, 2, 3, 4]
cs = Collection(1, 2, [3, Collection(4)])
assert cs == [1, 2, 3, 4]
# empty
cs = Collection()
assert cs == []
def test_collections_expand_to_indexes_with_pool():
cs = Collection(1, 2, Collection(3, 4))
cs = Collection(cs, pool=5)
assert cs == [1, 2, 3, 4]
cs = Collection(cs, pool=[4, 3, 2, 1])
assert cs == [3, 2, 1, 0]
cs = Collection(-Collection(1))
assert list(cs) == [-1]
cs.expand(pool=10)
assert list(cs) == [9]
cs = Collection(-Collection(1), pool=10)
assert list(cs) == [9]
def test_intersect():
# Intersect has to have 2 arguments
with pytest.raises(ValueError):
Intersect(1)
# without pool
left = Collection(1, 2, 3)
right = Collection(2, 3, 4)
ins = Intersect(left, right).expand()
assert list(ins) == [2, 3]
assert "Intersect" in repr(ins)
# indexes work
left = Collection(1, 2, 3, pool=list("abcd"))
right = Collection(2, 3, 4, pool=list("abcd"))
ins = Intersect(left, right).expand()
assert list(ins) == [2, 3] # returns 0-based indexes
# mixed str and index get correctly intersected
left = Collection(0, 1, "c", pool=list("abcd"))
right = Collection("b", 2, 3, pool=list("abcd"))
ins = Intersect(left, right).expand()
assert list(ins) == [1, 2]
def test_strings_turn_into_integers_with_string_pool():
cs = Collection(list("bcd"), pool=list("abc"))
assert cs == [1, 2]
assert cs.unmatched == {"d"}
assert "Collection" in repr(cs)
assert "[1, 2]" == str(cs)
def test_unmatched_saved():
cs = Collection(5, pool=4)
assert cs.unmatched == {5}
def test_collection_of_slice_expands_correctly():
cs = Collection(1, slice(3, 5))
assert cs == [1, 3, 4]
def test_inverted_collections_elements_get_unioned_then_inverted():
cs = Collection(~Collection(1, 2), ~Collection(3, 4), pool=5)
assert list(cs) == [0]
with pytest.raises(ValueError):
Collection(1, ~Collection(2)).expand(pool=3)