katylava
6/11/2012 - 10:36 PM

A smarter Python `slice` object.

A smarter Python slice object.

class SmartSlice(object):

    """
    A slice object which represents the mapping between old and new indices.

    Given a slice object, you can determine whether a given index will be
    present in the sub-list that slice represents. For example:

        >>> 3 in SmartSlice(2, 5)
        True
        >>> 0 in SmartSlice(2, 5)
        False
        >>> 5 in SmartSlice(2, 5)
        False

    This even works with a `step` argument:

        >>> 2 in SmartSlice(2, 10, 2)
        True
        >>> 4 in SmartSlice(2, 10, 2)
        True
        >>> 7 in SmartSlice(2, 10, 2)
        False

    You can map indices in the existing list to indices in the sub-list:

        >>> SmartSlice(2, 10, 2).forward_mapping(2)
        0
        >>> SmartSlice(2, 10, 2).forward_mapping(4)
        1
        >>> SmartSlice(2, 10, 2).forward_mapping(8)
        3

    If an index isn't present in the sub-list, an IndexError will be raised:

        >>> SmartSlice(2, 10, 2).forward_mapping(5)
        Traceback (most recent call last):
        ...
        IndexError: index 5 is not represented in the slice

    This works for reverse mapping as well (mapping indices in the slice to
    their original positions):

        >>> SmartSlice(2, 10, 2).reverse_mapping(0)
        2
        >>> SmartSlice(2, 10, 2).reverse_mapping(1)
        4
        >>> SmartSlice(2, 10, 2).reverse_mapping(3)
        8

    This may also raise an IndexError:

        >>> SmartSlice(2, 10, 2).reverse_mapping(12)
        Traceback (most recent call last):
        ...
        IndexError: the slice has no index 12

    If you provide a stop index, you can also reverse-map negative indices:

        >>> SmartSlice(2, 10, 2).reverse_mapping(-2)
        6

    And determine the length of the slice:

        >>> len(SmartSlice(2, 10, 2))
        4
        >>> len(SmartSlice(2, 11, 2))
        5

    But of course some things fail if you don't have a stop index, because
    there's no way of calculating them:

        >>> len(SmartSlice(2, None, 2))
        Traceback (most recent call last):
        ...
        TypeError: slice has no stop, and therefore no length
        >>> SmartSlice(2, None, 2).reverse_mapping(-2)
        Traceback (most recent call last):
        ...
        TypeError: slice has no stop, so does not support negative indexing
    """

    __slots__ = ('start', 'stop', 'step')

    def __init__(self, *args):
        if len(args) < 1 or len(args) > 3:
            raise TypeError("Expects between 1 and 3 arguments")
        elif len(args) == 1:
            start, stop, step = None, args[0], None
        elif len(args) == 2:
            start, stop, step = args[0], args[1], None
        elif len(args) == 3:
            start, stop, step = args[0], args[1], args[2]
        self.start = start
        self.stop = stop
        self.step = step

    def __contains__(self, i):
        if i < 0:
            raise TypeError("Negative indices are not supported for forward mapping")
        elif i < (self.start or 0):
            return False
        elif i >= (self.stop or float('inf')):
            return False
        return (i - (self.start or 0)) % (self.step or 1) == 0

    def __len__(self):
        if self.stop is None:
            raise TypeError("slice has no stop, and therefore no length")
        return ((self.stop + 1) - (self.start or 0)) // (self.step or 1)

    def forward_mapping(self, i):
        """Map indices in a list to their new indices in the slice."""

        if i not in self:
            raise IndexError("index %d is not represented in the slice" % i)
        return (i - (self.start or 0)) // (self.step or 1)

    def reverse_mapping(self, i):
        """Map indices in the slice to their old positions."""

        if (self.stop is not None) and i >= len(self) and i > 0:
            raise IndexError("the slice has no index %d" % i)
        elif i < 0:
            if self.stop is None:
                raise TypeError("slice has no stop, so does not support negative indexing")
            elif -i > len(self):
                raise IndexError("the slice has no index %d" % i)
            else:
                i = len(self) + i
        return (i * self.step) + (self.start or 0)