Skip to content
kmike edited this page Jul 21, 2012 · 3 revisions

A hack for adding pickle support to datrie.Trie (by marklr). Please note this is very inefficient; use trie.save()/datrie.Trie.load() if possible.

import cPickle, datrie, string
from datrie import AlphaMap

class PickledTrie(datrie.Trie):
    _alpha_map = None

    def __init__(self, alpha_map=None, _create=True):
        self._alpha_map = alpha_map
        super(PickledTrie, self).__init__(AlphaMap(alpha_map), _create)

    def __reduce__(self):
        return PickledTrie, (self._alpha_map, ), None, None, iter(self.items())

pt = PickledTrie(string.ascii_lowercase)
pt[u'foo'] = u'bar'

with open('tmp.dat', 'wb') as f:
    cPickle.dump(pt, f)

with open('tmp.dat', 'rb') as f:
    z = cPickle.load(f)
    print z.items()

Another variant (untested), more efficient but still hacky, by JustAMan:

import cPickle, datrie, tempfile, os

class PicklableTrie(datrie.Trie):
    def __reduce__(self):
        tempHandle, tempName = tempfile.mkstemp()
        try:
            os.close(tempHandle) # close the file we opened; now we'd save into it
            self.save(tempName)
            with open(tempName, 'rb') as inp:
                return datrie.__version__, inp.read()
        finally:
            os.unlink(tempName)

    def __setstate__(self, newState):
        assert datrie.__version__ == newState[0]
        tempHandle, tempName = tempfile.mkstemp()
        try:
            os.close(tempHandle) # close the file we opened; now we'd save into it
            with open(tempName, 'wb') as out:
                out.write(newState[1])
            self.load(tempName)
        finally:
            os.unlink(tempName)
Clone this wiki locally