defaultdict2 - a defaultdict whose default factory is called with the key

defaultdict2.py

from collections import defaultdict2
from unittest import TestCase

class defaultdict2(defaultdict):
    def __missing__(self, key):
        value = self.default_factory(key)
        self[key] = value
        return value

class TestDefaultdict2(TestCase):
    def test_defaultdict2(self):
        d = defaultdict2(lambda k: k + 1)

        # Assignment
        d[1] = 2
        d[3] = 5

        # Membership test
        self.assertTrue(1 in d)
        self.assertTrue(2 not in d)
        self.assertTrue(3 in d)

        # Subscripting
        self.assertEquals(d[1], 2)
        self.assertEquals(d[1], 2)
        self.assertEquals(d[2], 3)
        self.assertEquals(d[2], 3)
        self.assertEquals(d[3], 5)
        self.assertEquals(d[3], 5)

        # Deletion
        del d[2]
        del d[3]
        self.assertTrue(1 in d)
        self.assertTrue(2 not in d)
        self.assertTrue(3 not in d)

        # Still consistent after deletion
        self.assertEquals(d[1], 2)
        self.assertEquals(d[1], 2)
        self.assertEquals(d[2], 3)
        self.assertEquals(d[2], 3)
        self.assertEquals(d[3], 4)
        self.assertEquals(d[3], 4)

        # Dictiness
        self.assertEquals(d, {1:2, 2:3, 3:4})

        # Iteration
        for k, v in d.iteritems():
            self.assertEquals(v, k + 1)

if __name__ == "__main__":
    unittest.main()