diff --git a/patsy/test_highlevel.py b/patsy/test_highlevel.py index 7532d26..d809f6b 100644 --- a/patsy/test_highlevel.py +++ b/patsy/test_highlevel.py @@ -5,6 +5,7 @@ # Exhaustive end-to-end tests of the top-level API. import sys +from six.moves import cPickle as pickle import __future__ import numpy as np from nose.tools import assert_raises @@ -743,3 +744,17 @@ def test_C_and_pandas_categorical(): [[1, 0], [1, 1], [1, 0]]) + +def test_pickle_builder_roundtrips(): + design_matrix = dmatrix("x + a", {"x": [1, 2, 3], + "a": ["a1", "a2", "a3"]}) + builder = design_matrix.design_info.builder + + new_data = {"x": [10, 20, 30], + "a": ["a3", "a1", "a2"]} + m1 = dmatrix(builder, new_data) + + builder2 = pickle.loads(pickle.dumps(design_matrix.design_info.builder)) + m2 = dmatrix(builder2, new_data) + + assert np.allclose(m1, m2)