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)