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 

52 Certain URL schemes may require additional packages. For example, S3 

53 URLs require the ``s3fs`` library. See 

54 :ref:`install.optional_dependencies` for a full list. 

55 compression : str or dict, default 'infer' 

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

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

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

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

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

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

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

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

64 ``zipfile.ZipFile``, ``gzip.GzipFile``, 

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

66 ``tarfile.TarFile``, respectively. 

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

68 and to create a reproducible gzip archive: 

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

70 protocol : int 

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

72 default ``pickle.HIGHEST_PROTOCOL`` (see [1]_). A negative value for 

73 the protocol parameter is equivalent to setting its value to 

74 ``pickle.HIGHEST_PROTOCOL``. 

75 storage_options : dict, optional 

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

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

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

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

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

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

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

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

84 

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

86 

87 See Also 

88 -------- 

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

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

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

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

93 

94 Examples 

95 -------- 

96 >>> original_df = pd.DataFrame( 

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

98 ... ) # doctest: +SKIP 

99 >>> original_df # doctest: +SKIP 

100 foo bar 

101 0 0 5 

102 1 1 6 

103 2 2 7 

104 3 3 8 

105 4 4 9 

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

107 

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

109 >>> unpickled_df # doctest: +SKIP 

110 foo bar 

111 0 0 5 

112 1 1 6 

113 2 2 7 

114 3 3 8 

115 4 4 9 

116 """ 

117 if protocol < 0: 

118 protocol = pickle.HIGHEST_PROTOCOL 

119 

120 with get_handle( 

121 filepath_or_buffer, 

122 "wb", 

123 compression=compression, 

124 is_text=False, 

125 storage_options=storage_options, 

126 ) as handles: 

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

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

129 

130 

131@set_module("pandas") 

132def read_pickle( 

133 filepath_or_buffer: FilePath | ReadPickleBuffer, 

134 compression: CompressionOptions = "infer", 

135 storage_options: StorageOptions | None = None, 

136) -> DataFrame | Series: 

137 """ 

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

139 

140 .. warning:: 

141 

142 Loading pickled data received from untrusted sources can be 

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

144 

145 Parameters 

146 ---------- 

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

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

149 object implementing a binary ``read()`` function. 

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

151 

152 Certain URL schemes may require additional packages. For example, S3 

153 URLs require the ``s3fs`` library. See 

154 :ref:`install.optional_dependencies` for a full list. 

155 compression : str or dict, default 'infer' 

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

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

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

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

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

161 to be read in. 

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

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

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

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

166 ``zipfile.ZipFile``, ``gzip.GzipFile``, 

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

168 ``tarfile.TarFile``, respectively. 

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

170 using a custom compression dictionary: 

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

172 storage_options : dict, optional 

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

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

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

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

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

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

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

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

181 

182 Returns 

183 ------- 

184 object 

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

186 

187 See Also 

188 -------- 

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

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

191 read_hdf : Read HDF5 file into a DataFrame. 

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

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

194 

195 Notes 

196 ----- 

197 read_pickle is only guaranteed to be backwards compatible with pickles 

198 created by the current or previous major version of pandas, provided the 

199 object was serialized with to_pickle. For example, in pandas 3.x.y, the 

200 earliest supported pickle would be from 2.0.0. 

201 

202 Examples 

203 -------- 

204 >>> original_df = pd.DataFrame( 

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

206 ... ) # doctest: +SKIP 

207 >>> original_df # doctest: +SKIP 

208 foo bar 

209 0 0 5 

210 1 1 6 

211 2 2 7 

212 3 3 8 

213 4 4 9 

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

215 

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

217 >>> unpickled_df # doctest: +SKIP 

218 foo bar 

219 0 0 5 

220 1 1 6 

221 2 2 7 

222 3 3 8 

223 4 4 9 

224 """ 

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

226 # ValueError for legacy Timestamp pickles that mix a value with 

227 # by-component arguments (GH#31930) 

228 excs_to_catch = ( 

229 AttributeError, 

230 ImportError, 

231 ModuleNotFoundError, 

232 TypeError, 

233 ValueError, 

234 ) 

235 with get_handle( 

236 filepath_or_buffer, 

237 "rb", 

238 compression=compression, 

239 is_text=False, 

240 storage_options=storage_options, 

241 ) as handles: 

242 # 1) try standard library Pickle 

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

244 try: 

245 with warnings.catch_warnings(record=True): 

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

247 warnings.simplefilter("ignore", Warning) 

248 return pickle.load(handles.handle) 

249 except excs_to_catch: 

250 # e.g. 

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

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

253 handles.handle.seek(0) 

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