TST make sure DistMetric is picklable with both predefined as well as custom metric

raghavrv · raghavrv · commit dd043be5f32b · 2015-04-27T15:51:22.000+05:30
diff --git a/sklearn/neighbors/tests/test_dist_metrics.py b/sklearn/neighbors/tests/test_dist_metrics.py
@@ -1,4 +1,5 @@
 import itertools
+import pickle
 
 import numpy as np
 from numpy.testing import assert_array_almost_equal
@@ -9,6 +10,10 @@
 from nose import SkipTest
 
 
+def dist_func(x1, x2, p):
+    return np.sum((x1 - x2) ** p) ** (1. / p)
+
+
 def cmp_version(version1, version2):
     version1 = tuple(map(int, version1.split('.')[:2]))
     version2 = tuple(map(int, version2.split('.')[:2]))
@@ -122,18 +127,24 @@ def haversine_slow(x1, x2):
 
 
 def test_pyfunc_metric():
-    def dist_func(x1, x2, p):
-        return np.sum((x1 - x2) ** p) ** (1. / p)
-
     X = np.random.random((10, 3))
 
     euclidean = DistanceMetric.get_metric("euclidean")
     pyfunc = DistanceMetric.get_metric("pyfunc", func=dist_func, p=2)
 
+    # Check if both callable metric and predefined metric initialized
+    # DistanceMetric object is picklable
+    euclidean_pkl = pickle.loads(pickle.dumps(euclidean))
+    pyfunc_pkl = pickle.loads(pickle.dumps(pyfunc))
+
     D1 = euclidean.pairwise(X)
     D2 = pyfunc.pairwise(X)
 
+    D1_pkl = euclidean_pkl.pairwise(X)
+    D2_pkl = pyfunc_pkl.pairwise(X)
+
     assert_array_almost_equal(D1, D2)
+    assert_array_almost_equal(D1_pkl, D2_pkl)
 
 
 if __name__ == '__main__':
