10.26.2010

Being lazy in Python

One of my favorite features of Haskell is its lazy evaluation of list comprehensions. This makes it possible to take the nth position item of a list without first having to calculate the entire list. Pretty sweet if you want to reason of infinite data sets such as the Fibonacci numbers or primes. In a perfect world, I'd be writing Haskell all the time. However, that's just not possible. Python is far more popular, is installed by default on many OSes and can make some tasks very simple. So really, I want lazy list evaluations in Python. The only way to implement this without hacking up the actual interpreters object space is to mess with iterators. So I've implemented the Fibonacci sequence as an iterator. Essentially, I use the next parameter to move over the list of Fibonacci numbers and calculate the next one. I then store any numbers which I have already calculated to prevent extra work. I then use that pre-calculated list along with further calculation to implement list subscripting and contains. One thing to note is that itertools.takewhile will continue iteration for one through the failing condition, not up to. Hence the existence of a rollback function. Overall it was a fun exercise on using iterators and who knows, maybe it's actually useful.

from itertools import takewhile


class Fib(object):
    i = 0
    items = list()

    def __init__(self):
        self.n = 0
        self.n_1 = 1

    def __contains__(self, item):
        if item in self.items:
            return True
        else:
            self.items += list(takewhile(lambda x: x <= item, self))
            self._rollback()
            return self.items[-1]

    def __getitem__(self, key):
        if key <= len(self.items):
            return self.items[key - 1]
        else:
            self.items += list(takewhile(lambda x: self.i < key - 1, self))
            self._rollback()
            return self.items[-1]

    def __iter__(self):
        return self

    def next(self):
        next = self.n + self.n_1
        self.n = self.n_1
        self.n_1 = next
        self.i += 1
        return next

    def _rollback(self):
        """
        When iteration gets stopped by takewhile we've gone one too far
        so it needs to be backed up one
        """
        self.n = self.items[-2]
        self.n_1 = self.items[-1]
        self.i -= 1
Update: I've got a gist going of this which includes some bug fixes for slicing.