Search code examples
pythonclassslicepython-classobject-slicing

How to implement the lazy slicing of a class instance?


I am building a class which has underpinning it a list of ints, its __getitem__ method is then a reasonably computationally involved function outputting a list of potentially large size. Here is the code:

Class A(list):

    def __init__(self, n):
        self.li = [i for i in range(0, n)]
        super().__init__(self.li)
        self._ind = 0

    def __getitem__(self, item):

        assert (type(item) is int and item >= 0) or type(item) is slice

        if type(item) is slice:
            start, stop, step = item.start, item.stop, item.step
            if stop is None:
                stop = len(self)
            if step is not None:
                for i in range(start, stop, step):
                    yield self[i]
            else:
                for i in range(start, stop):
                    yield self[i]

        if item == 0 or item == 1:
            return []

        out = []
        while item != 1:
            if super().__getitem__(item):
                p = super().__getitem__(item)
            else:
                out.append(item)
                return out

            if p not in out:
                out.append(p)
            item //= p
        return out

My code is currently not working as it always returns a generator, and this messes up iterating over an instance of the class. However upon slicing, I do not wish it to do all the computations and store it, as this would make it consume a lot more memory. An example of what I wish to do:

from math import prod

test = A(10**7)
out = 0
for i in test[10**6:10**7]:
    out += prod(i)

How can I make slices work efficiently?


Solution

  • Instead of returning a generator, return a view. This is an object that conceptually represents the sequence of elements in question. We can do this by storing a reference to the original A instance, plus a range that encodes the instances. When we index into the view, we can ask the range which original index (or indices) are involved.

    Supposing we set up the general structure:

    class A(list):
        def __init__(self, n):
            super().__init__()
            self[:] = [i for i in range(0, n)]
    
        def _at(self, idx):
            # custom logic here to return the correct value for self[idx]
    
        def __getitem__(self, idx):
            if isinstance(idx, int):
                return self._at(idx)
            elif isinstance(idx, slice):
                # The `indices` method of a slice object converts to
                # start, stop, step values which we can use to construct a range.
                return A_view(self, range(*idx.indices(len(self))))
            else:
                raise TypeError # this is not an appropriate use for `assert`
    

    Then our view could look like:

    class A_view:
        def __init__(self, original, indices):
            self._original, self._indices = original, indices
    
    
        def __getitem__(self, idx):
            if isinstance(idx, int):
                return self._original[self._indices[idx]]
            elif isinstance(idx, slice):
                return A_view(self._original, self._indices[idx])
            else:
                raise TypeError
    
        def __len__(self):
            return len(self._indices)
    

    The idea is that if we receive an integer index, we translate it through the range object into an index for the original A, and call back to its __getitem__ (with an integer this time). If we receive another slice, we use it to slice our range into a sub-range, and make another view.

    Note that your A class should already be iterable, due to inheriting from list. To make the view iterable (and also get the in operator, forward and reverse iteration, .index and .count methods for free), you can inherit from collections.abc.Sequence. You just need a __getitem__ and __len__ - both easy to implement, as above - and the base class will do the rest (while also advertising to isinstance that your class has these properties).