Frozenset that caches state transitions

gistfile1.py

# Modeled after the idea used in dynamic language runtimes, where "expensive" mutable object transitions are cached, when you can safely assume that there will be a finite number of transitions.

# For example, on my iMac:

# $ python -m timeit -n 10000 -r 3 -s "from caching_frozenset import CachingFrozenset" "s = CachingFrozenset()
# for x in range(100):
#   s = s.with_value(x)"
# 10000 loops, best of 3: 66.1 usec per loop

# $ python -m timeit -n 10000 -r 3 "s = frozenset()
# for x in range(100):
#   s |= frozenset((x,))"
# 10000 loops, best of 3: 173 usec per loop

class CachingFrozenset(frozenset):
    # These dicts comprise the cache, and are shared by all instances
    _next_with = {}
    _next_without = {}

    @classmethod
    def reset_cache(cls):
        cls._next_with = {}
        cls._next_without = {}

    def with_value(self, new_value):
        if new_value in self:
            return self
        next_state = self._next_with.get((self, new_value), None)
        if next_state is None:
            next_state = self | frozenset([new_value])
            self._next_with[(self, new_value)] = next_state
        return next_state
    
    def without_value(self, old_value):
        if old_value not in self:
            return self
        next_state = self._next_without.get((self, old_value), None)
        if next_state is None:
            next_state = self - frozenset([old_value])
            self._next_without[(self, old_value)] = next_state
        return next_state

if __name__ == '__main__':
    import unittest

    class TestCachingFrozenset(unittest.TestCase):
        def setup(self):
            CachingFrozenset.reset_cache()

        def teardown(self):
            CachingFrozenset.reset_cache()

        def test_caching_frozenset(self):
            s = CachingFrozenset([1, 2, 3])
            assert 0 not in s
            assert 1 in s
            assert 2 in s
            assert 3 in s
            assert 4 not in s
            assert len(s) == 3
            assert list(sorted(s)) == [1, 2, 3]

        def test_with_value(self):
            s = CachingFrozenset()
            assert 1 not in s
            s_with_1 = s.with_value(1)
            assert 1 not in s
            assert 1 in s_with_1
            assert (s, 1) in CachingFrozenset._next_with
            assert CachingFrozenset._next_with[(s, 1)] is s_with_1

        def test_without_value(self):
            s = CachingFrozenset([1, 2, 3])
            assert 2 in s
            s_without_2 = s.without_value(2)
            assert 0 not in s_without_2
            assert 1 in s_without_2
            assert 2 not in s_without_2
            assert 3 in s_without_2
            assert 4 not in s_without_2
            assert 2 in s
            assert (s, 2) in CachingFrozenset._next_without
            assert CachingFrozenset._next_without[(s, 2)] is s_without_2

    unittest.main()