1
1
""" pickle compat """
2
2
import pickle
3
+ from typing import Any , Optional
3
4
import warnings
4
5
5
6
from pandas .compat import pickle_compat as pc
6
7
7
- from pandas .io . common import _get_handle , _stringify_path
8
+ from pandas ._typing import FilePathOrBuffer
8
9
10
+ from pandas .io .common import (
11
+ _get_handle ,
12
+ get_filepath_or_buffer as _get_filepath_or_buffer ,
13
+ )
9
14
10
- def to_pickle (obj , path , compression = "infer" , protocol = pickle .HIGHEST_PROTOCOL ):
15
+
16
+ def to_pickle (
17
+ obj : Any ,
18
+ filepath_or_buffer : FilePathOrBuffer ,
19
+ compression : Optional [str ] = "infer" ,
20
+ protocol : int = pickle .HIGHEST_PROTOCOL ,
21
+ ):
11
22
"""
12
23
Pickle (serialize) object to file.
13
24
@@ -16,10 +27,12 @@ def to_pickle(obj, path, compression="infer", protocol=pickle.HIGHEST_PROTOCOL):
16
27
obj : any object
17
28
Any python object.
18
29
path : str
19
- File path where the pickled object will be stored.
30
+ File path, URL, or buffer where the pickled object will be stored.
20
31
compression : {'infer', 'gzip', 'bz2', 'zip', 'xz', None}, default 'infer'
21
- A string representing the compression to use in the output file. By
22
- default, infers from the file extension in specified path.
32
+ If 'infer' and 'path_or_url' is path-like, then detect compression from
33
+ the following extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise no
34
+ compression) If 'infer' and 'path_or_url' is not path-like, then use
35
+ None (= no decompression).
23
36
protocol : int
24
37
Int which indicates which protocol should be used by the pickler,
25
38
default HIGHEST_PROTOCOL (see [1], paragraph 12.1.2). The possible
@@ -63,8 +76,12 @@ def to_pickle(obj, path, compression="infer", protocol=pickle.HIGHEST_PROTOCOL):
63
76
>>> import os
64
77
>>> os.remove("./dummy.pkl")
65
78
"""
66
- path = _stringify_path (path )
67
- f , fh = _get_handle (path , "wb" , compression = compression , is_text = False )
79
+ fp_or_buf , _ , compression , should_close = _get_filepath_or_buffer (
80
+ filepath_or_buffer , compression = compression , mode = "wb"
81
+ )
82
+ if not isinstance (fp_or_buf , str ) and compression == "infer" :
83
+ compression = None
84
+ f , fh = _get_handle (fp_or_buf , "wb" , compression = compression , is_text = False )
68
85
if protocol < 0 :
69
86
protocol = pickle .HIGHEST_PROTOCOL
70
87
try :
@@ -73,9 +90,16 @@ def to_pickle(obj, path, compression="infer", protocol=pickle.HIGHEST_PROTOCOL):
73
90
f .close ()
74
91
for _f in fh :
75
92
_f .close ()
93
+ if should_close :
94
+ try :
95
+ fp_or_buf .close ()
96
+ except ValueError :
97
+ pass
76
98
77
99
78
- def read_pickle (path , compression = "infer" ):
100
+ def read_pickle (
101
+ filepath_or_buffer : FilePathOrBuffer , compression : Optional [str ] = "infer"
102
+ ):
79
103
"""
80
104
Load pickled pandas object (or any object) from file.
81
105
@@ -86,13 +110,13 @@ def read_pickle(path, compression="infer"):
86
110
87
111
Parameters
88
112
----------
89
- path : str
90
- File path where the pickled object will be loaded.
113
+ filepath_or_buffer : str
114
+ File path, URL, or buffer where the pickled object will be loaded from .
91
115
compression : {'infer', 'gzip', 'bz2', 'zip', 'xz', None}, default 'infer'
92
- For on-the-fly decompression of on-disk data. If 'infer', then use
93
- gzip, bz2, xz or zip if path ends in '.gz', '.bz2', '.xz',
94
- or '.zip' respectively, and no decompression otherwise.
95
- Set to None for no decompression.
116
+ If 'infer' and 'path_or_url' is path-like , then detect compression from
117
+ the following extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise no
118
+ compression) If 'infer' and 'path_or_url' is not path-like, then use
119
+ None (= no decompression) .
96
120
97
121
Returns
98
122
-------
@@ -134,8 +158,12 @@ def read_pickle(path, compression="infer"):
134
158
>>> import os
135
159
>>> os.remove("./dummy.pkl")
136
160
"""
137
- path = _stringify_path (path )
138
- f , fh = _get_handle (path , "rb" , compression = compression , is_text = False )
161
+ fp_or_buf , _ , compression , should_close = _get_filepath_or_buffer (
162
+ filepath_or_buffer , compression = compression
163
+ )
164
+ if not isinstance (fp_or_buf , str ) and compression == "infer" :
165
+ compression = None
166
+ f , fh = _get_handle (fp_or_buf , "rb" , compression = compression , is_text = False )
139
167
140
168
# 1) try standard library Pickle
141
169
# 2) try pickle_compat (older pandas version) to handle subclass changes
@@ -159,3 +187,8 @@ def read_pickle(path, compression="infer"):
159
187
f .close ()
160
188
for _f in fh :
161
189
_f .close ()
190
+ if should_close :
191
+ try :
192
+ fp_or_buf .close ()
193
+ except ValueError :
194
+ pass
0 commit comments