forked from cloudpipe/cloudpickle
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cloudpickle_file_test.py
120 lines (102 loc) · 4.16 KB
/
cloudpickle_file_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
import unittest
import tempfile
import os
import shutil
import pickle
import sys
try:
from io import StringIO
except ImportError:
# compat for Python 2.6
from StringIO import StringIO
import pytest
from mock import patch, mock_open
import cloudpickle
class CloudPickleFileTests(unittest.TestCase):
"""In Cloudpickle, expected behaviour when pickling an opened file
is to send its contents over the wire and seek to the same position."""
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.tmpfilepath = os.path.join(self.tmpdir, 'testfile')
self.teststring = u'Hello world!'
def tearDown(self):
shutil.rmtree(self.tmpdir)
def test_empty_file(self):
# Empty file
open(self.tmpfilepath, 'w').close()
with open(self.tmpfilepath, 'r') as f:
self.assertEquals('', pickle.loads(cloudpickle.dumps(f)).read())
os.remove(self.tmpfilepath)
def test_closed_file(self):
# Write & close
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
with pytest.raises(pickle.PicklingError) as excinfo:
cloudpickle.dumps(f)
assert "Cannot pickle closed files" in str(excinfo.value)
os.remove(self.tmpfilepath)
def test_r_mode(self):
# Write & close
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
# Open for reading
with open(self.tmpfilepath, 'r') as f:
new_f = pickle.loads(cloudpickle.dumps(f))
self.assertEquals(self.teststring, new_f.read())
os.remove(self.tmpfilepath)
def test_w_mode(self):
with open(self.tmpfilepath, 'w') as f:
f.write(self.teststring)
f.seek(0)
self.assertRaises(pickle.PicklingError,
lambda: cloudpickle.dumps(f))
os.remove(self.tmpfilepath)
def test_plus_mode(self):
# Write, then seek to 0
with open(self.tmpfilepath, 'w+') as f:
f.write(self.teststring)
f.seek(0)
new_f = pickle.loads(cloudpickle.dumps(f))
self.assertEquals(self.teststring, new_f.read())
os.remove(self.tmpfilepath)
def test_seek(self):
# Write, then seek to arbitrary position
with open(self.tmpfilepath, 'w+') as f:
f.write(self.teststring)
f.seek(4)
unpickled = pickle.loads(cloudpickle.dumps(f))
# unpickled StringIO is at position 4
self.assertEquals(4, unpickled.tell())
self.assertEquals(self.teststring[4:], unpickled.read())
# but unpickled StringIO also contained the start
unpickled.seek(0)
self.assertEquals(self.teststring, unpickled.read())
os.remove(self.tmpfilepath)
@pytest.mark.skipif(sys.version_info > (2, 7),
reason="only works on Python 2.x")
def test_temp_file(self):
with tempfile.NamedTemporaryFile(mode='ab+') as fp:
fp.write(self.teststring.encode('UTF-8'))
fp.seek(0)
f = fp.file
# FIXME this doesn't work yet: cloudpickle.dumps(fp)
newfile = pickle.loads(cloudpickle.dumps(f))
self.assertEquals(self.teststring, newfile.read())
def test_pickling_special_file_handles(self):
# Warning: if you want to run your tests with nose, add -s option
for out in sys.stdout, sys.stderr: # Regression test for SPARK-3415
self.assertEquals(out, pickle.loads(cloudpickle.dumps(out)))
self.assertRaises(pickle.PicklingError,
lambda: cloudpickle.dumps(sys.stdin))
def NOT_WORKING_test_tty(self):
# FIXME: Mocking 'file' is not trivial... and fails for now
from sys import version_info
if version_info.major == 2:
import __builtin__ as builtins # pylint:disable=import-error
else:
import builtins # pylint:disable=import-error
with patch.object(builtins, 'open', mock_open(), create=True):
with open('foo', 'w+') as handle:
cloudpickle.dumps(handle)
if __name__ == '__main__':
unittest.main()