Saving KDTree object in Python?

KDtree uses nested classes to define its node types (innernode, leafnode). Pickle only works on module-level class definitions, so a nested class trips it up:

import cPickleclass Foo(object):    class Bar(object):        passobj = Foo.Bar()print obj.__class__cPickle.dumps(obj)<class '__main__.Bar'>cPickle.PicklingError: Can't pickle <class '__main__.Bar'>: attribute lookup __main__.Bar failed

However, there is a (hacky) workaround by monkey-patching the class definitions into the scipy.spatial.kdtree at module scope so the pickler can find them. If all of your code which reads and writes pickled KDtree objects installs these patches, this hack should work fine:

import cPickleimport numpyfrom scipy.spatial import kdtree# patch module-level attribute to enable pickle to workkdtree.node = kdtree.KDTree.nodekdtree.leafnode = kdtree.KDTree.leafnodekdtree.innernode = kdtree.KDTree.innernodex, y = numpy.mgrid[0:5, 2:8]t1 = kdtree.KDTree(zip(x.ravel(), y.ravel()))r1 = t1.query([3.4, 4.1])raw = cPickle.dumps(t1)# read in the pickled treet2 = cPickle.loads(raw)r2 = t2.query([3.4, 4.1])print t1.tree.__class__print repr(raw)[:70]print t1.data[r1[1]], t2.data[r2[1]]

Output:

<class 'scipy.spatial.kdtree.innernode'>"ccopy_reg\n_reconstructor\np1\n(cscipy.spatial.kdtree\nKDTree\np2\nc_[3 4] [3 4]