From c53dd6df18e898e4277f965ad06497bdca5856ea Mon Sep 17 00:00:00 2001 From: Hoyt Koepke Date: Tue, 29 Sep 2020 15:41:59 -0600 Subject: [PATCH] Update to SFrame serialization to fix issues with python 2.7 and cloudpickle. (#3333) * Update to fix issue with python 2.7 and cloudpickle. * Addressed PR feedback. Co-authored-by: Hoyt Koepke --- src/python/turicreate/data_structures/serialization.py | 8 +++++++- src/python/turicreate/test/test_sframe.py | 7 +++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/python/turicreate/data_structures/serialization.py b/src/python/turicreate/data_structures/serialization.py index b580ee2eaa..efe6bcca36 100644 --- a/src/python/turicreate/data_structures/serialization.py +++ b/src/python/turicreate/data_structures/serialization.py @@ -8,6 +8,7 @@ # SFrames require a disk-backed file or directory to work with. This directory # has to be present to allow for serialization or deserialization. __serialization_directory = None +from sys import version_info as _version_info def enable_sframe_serialization(serialization_directory): @@ -28,6 +29,8 @@ def enable_sframe_serialization(serialization_directory): import os + if _version_info[0] == 2: + raise RuntimeError("SFrame serialization not supported in python 2.") global __serialization_directory if serialization_directory is None: __serialization_directory = None @@ -61,7 +64,10 @@ def _safe_serialization_directory(): from pickle import PickleError if __serialization_directory is None: - raise PickleError("Serialization directory not set to enable pickling or unpickling. " + + expected_error = TypeError if (_version_info[0] == 3) else PickleError + + raise expected_error("Serialization directory not set to enable pickling or unpickling. " "Set using turicreate.data_structures.serialization.enable_sframe_serialization().") return __serialization_directory diff --git a/src/python/turicreate/test/test_sframe.py b/src/python/turicreate/test/test_sframe.py index ef706eb824..9b3633e4c3 100644 --- a/src/python/turicreate/test/test_sframe.py +++ b/src/python/turicreate/test/test_sframe.py @@ -12,6 +12,7 @@ from ..util import _assert_sframe_equal, generate_random_sframe from .. import _launch, load_sframe, aggregate from . import util +from sys import version_info import pandas as pd from .._cython.cy_flexible_type import GMT @@ -642,6 +643,7 @@ def test_save_to_csv(self): f.close() os.unlink(f.name) + @pytest.mark.skipif(version_info[0] != 3, "Not a supported feature in python 2.7") def test_pickling(self): import pickle @@ -650,9 +652,10 @@ def test_pickling(self): X = generate_random_sframe(100, "ncc") with util.TempDirectory() as f: - - self.assertRaises(pickle.PickleError, lambda: pickle.dumps(X)) + expected_error = TypeError if (version_info[0] == 3) else PicklingError + + self.assertRaises(expected_error, lambda: pickle.dumps(X)) serialization.enable_sframe_serialization(f)