Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/pickle.py: 46%

1"""pickle compat""" 

2 

3from __future__ import annotations 

4 

5import pickle 

6from typing import ( 

7 TYPE_CHECKING, 

8 Any, 

9) 

10import warnings 

11 

12from pandas.compat import pickle_compat 

13from pandas.util._decorators import set_module 

14 

15from pandas.io.common import get_handle 

16 

17if TYPE_CHECKING: 

18 from pandas._typing import ( 

19 CompressionOptions, 

20 FilePath, 

21 ReadPickleBuffer, 

22 StorageOptions, 

23 WriteBuffer, 

24 ) 

25 

26 from pandas import ( 

27 DataFrame, 

28 Series, 

29 ) 

30 

31 

32@set_module("pandas") 

33def to_pickle( 

34 obj: Any, 

35 filepath_or_buffer: FilePath | WriteBuffer[bytes], 

36 compression: CompressionOptions = "infer", 

37 protocol: int = pickle.HIGHEST_PROTOCOL, 

38 storage_options: StorageOptions | None = None, 

39) -> None: 

40 """ 

41 Pickle (serialize) object to file. 

42 

43 Parameters 

44 ---------- 

45 obj : any object 

46 Any python object. 

47 filepath_or_buffer : str, path object, or file-like object 

48 String, path object (implementing ``os.PathLike[str]``), or file-like 

49 object implementing a binary ``write()`` function. 

50 Also accepts URL. URL has to be of S3 or GCS. 

51 compression : str or dict, default 'infer' 

52 For on-the-fly compression of the output data. If 'infer' and 

53 'filepath_or_buffer' is path-like, then detect compression from the 

54 following extensions: '.gz', '.bz2', '.zip', '.xz', '.zst', '.tar', 

55 '.tar.gz', '.tar.xz' or '.tar.bz2' (otherwise no compression). 

56 Set to ``None`` for no compression. 

57 Can also be a dict with key ``'method'`` set 

58 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, 

59 ``'tar'``} and other key-value pairs are forwarded to 

60 ``zipfile.ZipFile``, ``gzip.GzipFile``, 

61 ``bz2.BZ2File``, ``zstandard.ZstdCompressor``, ``lzma.LZMAFile`` or 

62 ``tarfile.TarFile``, respectively. 

63 As an example, the following could be passed for faster compression 

64 and to create a reproducible gzip archive: 

65 ``compression={'method': 'gzip', 'compresslevel': 1, 'mtime': 1}``. 

66 protocol : int 

67 Int which indicates which protocol should be used by the pickler, 

68 default HIGHEST_PROTOCOL (see [1], paragraph 12.1.2). The possible 

69 values for this parameter depend on the version of Python. For Python 

70 2.x, possible values are 0, 1, 2. For Python>=3.0, 3 is a valid value. 

71 For Python >= 3.4, 4 is a valid value. A negative value for the 

72 protocol parameter is equivalent to setting its value to 

73 HIGHEST_PROTOCOL. 

74 storage_options : dict, optional 

75 Extra options that make sense for a particular storage connection, e.g. 

76 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs 

77 are forwarded to ``urllib.request.Request`` as header options. For other 

78 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are 

79 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more 

80 details, and for more examples on storage options refer `here 

81 <https://pandas.pydata.org/docs/user_guide/io.html? 

82 highlight=storage_options#reading-writing-remote-files>`_. 

83 

84 .. [1] https://docs.python.org/3/library/pickle.html 

85 

86 See Also 

87 -------- 

88 read_pickle : Load pickled pandas object (or any object) from file. 

89 DataFrame.to_hdf : Write DataFrame to an HDF5 file. 

90 DataFrame.to_sql : Write DataFrame to a SQL database. 

91 DataFrame.to_parquet : Write a DataFrame to the binary parquet format. 

92 

93 Examples 

94 -------- 

95 >>> original_df = pd.DataFrame( 

96 ... {"foo": range(5), "bar": range(5, 10)} 

97 ... ) # doctest: +SKIP 

98 >>> original_df # doctest: +SKIP 

99 foo bar 

100 0 0 5 

101 1 1 6 

102 2 2 7 

103 3 3 8 

104 4 4 9 

105 >>> pd.to_pickle(original_df, "./dummy.pkl") # doctest: +SKIP 

106 

107 >>> unpickled_df = pd.read_pickle("./dummy.pkl") # doctest: +SKIP 

108 >>> unpickled_df # doctest: +SKIP 

109 foo bar 

110 0 0 5 

111 1 1 6 

112 2 2 7 

113 3 3 8 

114 4 4 9 

115 """ 

116 if protocol < 0: 

117 protocol = pickle.HIGHEST_PROTOCOL 

118 

119 with get_handle( 

120 filepath_or_buffer, 

121 "wb", 

122 compression=compression, 

123 is_text=False, 

124 storage_options=storage_options, 

125 ) as handles: 

126 # letting pickle write directly to the buffer is more memory-efficient 

127 pickle.dump(obj, handles.handle, protocol=protocol) 

128 

129 

130@set_module("pandas") 

131def read_pickle( 

132 filepath_or_buffer: FilePath | ReadPickleBuffer, 

133 compression: CompressionOptions = "infer", 

134 storage_options: StorageOptions | None = None, 

135) -> DataFrame | Series: 

136 """ 

137 Load pickled pandas object (or any object) from file and return unpickled object. 

138 

139 .. warning:: 

140 

141 Loading pickled data received from untrusted sources can be 

142 unsafe. See `here <https://docs.python.org/3/library/pickle.html>`__. 

143 

144 Parameters 

145 ---------- 

146 filepath_or_buffer : str, path object, or file-like object 

147 String, path object (implementing ``os.PathLike[str]``), or file-like 

148 object implementing a binary ``readlines()`` function. 

149 Also accepts URL. URL is not limited to S3 and GCS. 

150 compression : str or dict, default 'infer' 

151 For on-the-fly decompression of on-disk data. If 'infer' and 

152 'filepath_or_buffer' is path-like, then detect compression from the 

153 following extensions: '.gz', '.bz2', '.zip', '.xz', '.zst', '.tar', 

154 '.tar.gz', '.tar.xz' or '.tar.bz2' (otherwise no compression). 

155 If using 'zip' or 'tar', the ZIP file must contain only one data file 

156 to be read in. 

157 Set to ``None`` for no decompression. 

158 Can also be a dict with key ``'method'`` set 

159 to one of {``'zip'``, ``'gzip'``, ``'bz2'``, ``'zstd'``, ``'xz'``, 

160 ``'tar'``} and other key-value pairs are forwarded to 

161 ``zipfile.ZipFile``, ``gzip.GzipFile``, 

162 ``bz2.BZ2File``, ``zstandard.ZstdDecompressor``, ``lzma.LZMAFile`` or 

163 ``tarfile.TarFile``, respectively. 

164 As an example, the following could be passed for Zstandard decompression 

165 using a custom compression dictionary: 

166 ``compression={'method': 'zstd', 'dict_data': my_compression_dict}``. 

167 storage_options : dict, optional 

168 Extra options that make sense for a particular storage connection, e.g. 

169 host, port, username, password, etc. For HTTP(S) URLs the key-value pairs 

170 are forwarded to ``urllib.request.Request`` as header options. For other 

171 URLs (e.g. starting with "s3://", and "gcs://") the key-value pairs are 

172 forwarded to ``fsspec.open``. Please see ``fsspec`` and ``urllib`` for more 

173 details, and for more examples on storage options refer `here 

174 <https://pandas.pydata.org/docs/user_guide/io.html? 

175 highlight=storage_options#reading-writing-remote-files>`_. 

176 

177 Returns 

178 ------- 

179 object 

180 The unpickled pandas object (or any object) that was stored in file. 

181 

182 See Also 

183 -------- 

184 DataFrame.to_pickle : Pickle (serialize) DataFrame object to file. 

185 Series.to_pickle : Pickle (serialize) Series object to file. 

186 read_hdf : Read HDF5 file into a DataFrame. 

187 read_sql : Read SQL query or database table into a DataFrame. 

188 read_parquet : Load a parquet object, returning a DataFrame. 

189 

190 Notes 

191 ----- 

192 read_pickle is only guaranteed to be backwards compatible to pandas 1.0 

193 provided the object was serialized with to_pickle. 

194 

195 Examples 

196 -------- 

197 >>> original_df = pd.DataFrame( 

198 ... {"foo": range(5), "bar": range(5, 10)} 

199 ... ) # doctest: +SKIP 

200 >>> original_df # doctest: +SKIP 

201 foo bar 

202 0 0 5 

203 1 1 6 

204 2 2 7 

205 3 3 8 

206 4 4 9 

207 >>> pd.to_pickle(original_df, "./dummy.pkl") # doctest: +SKIP 

208 

209 >>> unpickled_df = pd.read_pickle("./dummy.pkl") # doctest: +SKIP 

210 >>> unpickled_df # doctest: +SKIP 

211 foo bar 

212 0 0 5 

213 1 1 6 

214 2 2 7 

215 3 3 8 

216 4 4 9 

217 """ 

218 # TypeError for Cython complaints about object.__new__ vs Tick.__new__ 

219 excs_to_catch = (AttributeError, ImportError, ModuleNotFoundError, TypeError) 

220 with get_handle( 

221 filepath_or_buffer, 

222 "rb", 

223 compression=compression, 

224 is_text=False, 

225 storage_options=storage_options, 

226 ) as handles: 

227 # 1) try standard library Pickle 

228 # 2) try pickle_compat (older pandas version) to handle subclass changes 

229 try: 

230 with warnings.catch_warnings(record=True): 

231 # We want to silence any warnings about, e.g. moved modules. 

232 warnings.simplefilter("ignore", Warning) 

233 return pickle.load(handles.handle) 

234 except excs_to_catch: 

235 # e.g. 

236 # "No module named 'pandas.core.sparse.series'" 

237 # "Can't get attribute '_nat_unpickle' on <module 'pandas._libs.tslib" 

238 handles.handle.seek(0) 

239 return pickle_compat.Unpickler(handles.handle).load()