Merge pull request #81 from dato-code/gl_pickle_fix

GLPickle expandvars and expanduser paths are now supported.

Author: Hoyt Koepke

oss_src/unity/python/sframe/_gl_pickle.py

@@ -252,7 +252,9 @@ def __init__(self, filename, protocol = -1, min_bytes_to_save = 0):
             self.hadoop_conf_dir = None
         else:
             # Make sure the directory exists.
-            filename = _os.path.abspath(filename)
+            filename = _os.path.abspath(
+                         _os.path.expanduser(
+                           _os.path.expandvars(filename)))
             if not _os.path.exists(filename):
                 _os.makedirs(filename)
             elif _os.path.isdir(filename):
@@ -457,6 +459,9 @@ def __init__(self, filename):
             _file_util.download_from_hdfs(filename, self.tmp_file)
             filename = self.tmp_file
         else:
+            filename = _os.path.abspath(
+                         _os.path.expanduser(
+                           _os.path.expandvars(filename)))
             if not _os.path.exists(filename):
                 raise IOError('%s is not a valid file name.' % filename)
 

oss_src/unity/python/sframe/test/test_gl_pickler.py

@@ -40,7 +40,7 @@ def test_pickling_simple_types(self):
             pickler.dump(obj)
             pickler.close()
             obj_ret = gl_pickle.GLUnpickler(self.filename).load()
-            assert obj == obj_ret, "Failed pickling in %s (Got back %s)" % (obj, obj_ret)
+            self.assertEqual(obj, obj_ret)
 
     def test_pickling_sarray_types(self):
 
@@ -54,8 +54,7 @@ def test_pickling_sarray_types(self):
             pickler.dump(obj)
             pickler.close()
             obj_ret = gl_pickle.GLUnpickler(self.filename).load()
-            assert list(obj) ==  list(obj_ret), \
-                       "Failed pickling in %s (Got back %s)" % (obj, obj_ret)
+            self.assertEqual(list(obj), list(obj_ret))
 
     def test_pickling_sframe_types(self):
 
@@ -120,7 +119,7 @@ def test_combination_gl_python_types(self):
             assert_sframe_equal(obj[0].get_vertices(), obj_ret[0].get_vertices())
             assert_sframe_equal(obj[0].get_edges(), obj_ret[0].get_edges())
             assert_sframe_equal(obj[1], obj_ret[1])
-            assert list(obj[2]) == list(obj_ret[2])
+            self.assertEqual(list(obj[2]), list(obj_ret[2]))
 
     def test_pickle_compatibility(self):
         obj_list = [
@@ -135,8 +134,7 @@ def test_pickle_compatibility(self):
             pickler.dump(obj)
             file.close()
             obj_ret = gl_pickle.GLUnpickler(self.filename).load()
-            assert obj == obj_ret, \
-                "Failed pickling in %s (Got back %s)" % (obj, obj_ret)
+            self.assertTrue(obj, obj_ret)
 
     def test_cloud_pickle_compatibility(self):
         obj_list = [
@@ -151,8 +149,7 @@ def test_cloud_pickle_compatibility(self):
             pickler.dump(obj)
             file.close()
             obj_ret = gl_pickle.GLUnpickler(self.filename).load()
-            assert obj == obj_ret, \
-                "Failed pickling in %s (Got back %s)" % (obj, obj_ret)
+            self.assertTrue(obj, obj_ret)
 
     def test_relative_path(self):
         # Arrange
@@ -171,6 +168,47 @@ def test_relative_path(self):
         # Clean up
         shutil.rmtree(relative_path)
 
+    def test_user_path(self):
+        # Arrange
+        sf1 = SFrame(range(10))
+        relative_path = '~/tmp/%s' % self.filename
+
+
+        # Act
+        pickler = gl_pickle.GLPickler(relative_path)
+        pickler.dump(sf1)
+        pickler.close()
+        sf2 = gl_pickle.GLUnpickler(relative_path).load()
+
+        # Assert
+        abs_path = os.path.expanduser(relative_path)
+        self.assertTrue(os.path.exists(abs_path))
+        assert_sframe_equal(sf1, sf2)
+
+        # Clean up
+        shutil.rmtree(abs_path)
+
+    def test_env_path(self):
+        # Arrange
+        sf1 = SFrame(range(10))
+        os.environ["FOO"] = '~/bar'
+        relative_path = '${FOO}/tmp/%s' % self.filename
+
+        # Act
+        pickler = gl_pickle.GLPickler(relative_path)
+        pickler.dump(sf1)
+        pickler.close()
+        sf2 = gl_pickle.GLUnpickler(relative_path).load()
+
+        # Assert
+        abs_path = os.path.expanduser(os.path.expandvars(relative_path))
+        self.assertTrue(os.path.exists(abs_path))
+        assert_sframe_equal(sf1, sf2)
+
+        # Clean up
+        shutil.rmtree(abs_path)
+        del(os.environ["FOO"])
+
     @moto.mock_s3
     def test_save_to_s3(self):
         # Arrange