Reduce I/O by fusing repeated getitems on arrays about dask HOT 7 CLOSED

from operator import getitem

def test_fuse_getitem():
    pairs = [((getitem, (getitem, 'x', slice(1000, 2000)), slice(15, 20)),
              (getitem, 'x', slice(1015, 1020))),

             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(100, 200))),
                        (slice(15, 20), slice(50, 60))),
              (getitem, 'x', (slice(1015, 1020), slice(150, 160)))),

             ((getitem, (getitem, 'x', slice(1000, 2000)), 10),
              (getitem, 'x', 1010)),

             ((getitem, (getitem, 'x', (slice(1000, 2000), 10)),
                        (slice(15, 20),)),
              (getitem, 'x', (slice(1015, 1020), 10))),

             ((getitem, (getitem, 'x', (10, slice(1000, 2000))),
                        (slice(15, 20),)),
              (getitem, 'x', (10, slice(1015, 1020))))

        ]

    for inp, expected in pairs:
        result = fuse_getitem(inp)
        assert result == expected


def fuse_getitem(term):
    (_, (_, x, a), b) = term

    c = fuse_slice(a, b)

    return (getitem, x, c)


def normalize_slice(s):
    start, stop, step = s.start, s.stop, s.step
    if start is None:
        start = 0
    if step is None:
        step = 1
    return slice(start, stop, step)

def fuse_slice(a, b):
    """

    >>> fuse_slice(slice(1000, 2000), slice(10, 15))
    slice(1010, 1015, None)

    >>> fuse_slice(slice(1000, 2000), 10)
    1010

    >>> fuse_slice(slice(1000, 2000, 5), 10)
    1050

    >>> fuse_slice(slice(1000, 2000, 5), slice(10, 20, 2))
    slice(1050, 1100, 10)
    """

    if isinstance(a, slice):
        a = normalize_slice(a)
    if isinstance(b, slice):
        b = normalize_slice(b)

    if isinstance(a, slice) and isinstance(b, int):
        return a.start + b*(a.step if a.step is not None else 1)
    if isinstance(a, slice) and isinstance(b, slice):
        start = a.start + a.step * b.start
        if b.stop is not None:
            stop = a.start + a.step * b.stop
        else:
            stop = None
        if a.stop is not None:
            if stop is not None:
                stop = min(a.stop, stop)
            else:
                stop = a.stop
            stop = stop - ((stop - start) % b.step)
        step = a.step * b.step
        if step == 1:
            step = None
        return slice(start, stop, step)
    if isinstance(a, tuple) and not isinstance(b, tuple):
        b = (b,)
    if isinstance(a, tuple) and isinstance(b, tuple):
        j = 0
        result = list()
        for i in range(len(a)):
            if isinstance(a[i], int):
                result.append(a[i])
                continue
            result.append(fuse_slice(a[i], b[j]))
            j += 1
        return tuple(result)
    raise NotImplementedError()