Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/pandas/compat/pickle_compat.py: 29%

1""" 

2Support pre-0.12 series pickle compatibility. 

3""" 

4from __future__ import annotations 

5 

6import contextlib 

7import copy 

8import io 

9import pickle as pkl 

10from typing import TYPE_CHECKING 

11 

12import numpy as np 

13 

14from pandas._libs.arrays import NDArrayBacked 

15from pandas._libs.tslibs import BaseOffset 

16 

17from pandas import Index 

18from pandas.core.arrays import ( 

19 DatetimeArray, 

20 PeriodArray, 

21 TimedeltaArray, 

22) 

23from pandas.core.internals import BlockManager 

24 

25if TYPE_CHECKING: 

26 from collections.abc import Generator 

27 

28 

29def load_reduce(self) -> None: 

30 stack = self.stack 

31 args = stack.pop() 

32 func = stack[-1] 

33 

34 try: 

35 stack[-1] = func(*args) 

36 return 

37 except TypeError as err: 

38 # If we have a deprecated function, 

39 # try to replace and try again. 

40 

41 msg = "_reconstruct: First argument must be a sub-type of ndarray" 

42 

43 if msg in str(err): 

44 try: 

45 cls = args[0] 

46 stack[-1] = object.__new__(cls) 

47 return 

48 except TypeError: 

49 pass 

50 elif args and isinstance(args[0], type) and issubclass(args[0], BaseOffset): 

51 # TypeError: object.__new__(Day) is not safe, use Day.__new__() 

52 cls = args[0] 

53 stack[-1] = cls.__new__(*args) 

54 return 

55 elif args and issubclass(args[0], PeriodArray): 

56 cls = args[0] 

57 stack[-1] = NDArrayBacked.__new__(*args) 

58 return 

59 

60 raise 

61 

62 

63# If classes are moved, provide compat here. 

64_class_locations_map = { 

65 ("pandas.core.sparse.array", "SparseArray"): ("pandas.core.arrays", "SparseArray"), 

66 # 15477 

67 ("pandas.core.base", "FrozenNDArray"): ("numpy", "ndarray"), 

68 # Re-routing unpickle block logic to go through _unpickle_block instead 

69 # for pandas <= 1.3.5 

70 ("pandas.core.internals.blocks", "new_block"): ( 

71 "pandas._libs.internals", 

72 "_unpickle_block", 

73 ), 

74 ("pandas.core.indexes.frozen", "FrozenNDArray"): ("numpy", "ndarray"), 

75 ("pandas.core.base", "FrozenList"): ("pandas.core.indexes.frozen", "FrozenList"), 

76 # 10890 

77 ("pandas.core.series", "TimeSeries"): ("pandas.core.series", "Series"), 

78 ("pandas.sparse.series", "SparseTimeSeries"): ( 

79 "pandas.core.sparse.series", 

80 "SparseSeries", 

81 ), 

82 # 12588, extensions moving 

83 ("pandas._sparse", "BlockIndex"): ("pandas._libs.sparse", "BlockIndex"), 

84 ("pandas.tslib", "Timestamp"): ("pandas._libs.tslib", "Timestamp"), 

85 # 18543 moving period 

86 ("pandas._period", "Period"): ("pandas._libs.tslibs.period", "Period"), 

87 ("pandas._libs.period", "Period"): ("pandas._libs.tslibs.period", "Period"), 

88 # 18014 moved __nat_unpickle from _libs.tslib-->_libs.tslibs.nattype 

89 ("pandas.tslib", "__nat_unpickle"): ( 

90 "pandas._libs.tslibs.nattype", 

91 "__nat_unpickle", 

92 ), 

93 ("pandas._libs.tslib", "__nat_unpickle"): ( 

94 "pandas._libs.tslibs.nattype", 

95 "__nat_unpickle", 

96 ), 

97 # 15998 top-level dirs moving 

98 ("pandas.sparse.array", "SparseArray"): ( 

99 "pandas.core.arrays.sparse", 

100 "SparseArray", 

101 ), 

102 ("pandas.indexes.base", "_new_Index"): ("pandas.core.indexes.base", "_new_Index"), 

103 ("pandas.indexes.base", "Index"): ("pandas.core.indexes.base", "Index"), 

104 ("pandas.indexes.numeric", "Int64Index"): ( 

105 "pandas.core.indexes.base", 

106 "Index", # updated in 50775 

107 ), 

108 ("pandas.indexes.range", "RangeIndex"): ("pandas.core.indexes.range", "RangeIndex"), 

109 ("pandas.indexes.multi", "MultiIndex"): ("pandas.core.indexes.multi", "MultiIndex"), 

110 ("pandas.tseries.index", "_new_DatetimeIndex"): ( 

111 "pandas.core.indexes.datetimes", 

112 "_new_DatetimeIndex", 

113 ), 

114 ("pandas.tseries.index", "DatetimeIndex"): ( 

115 "pandas.core.indexes.datetimes", 

116 "DatetimeIndex", 

117 ), 

118 ("pandas.tseries.period", "PeriodIndex"): ( 

119 "pandas.core.indexes.period", 

120 "PeriodIndex", 

121 ), 

122 # 19269, arrays moving 

123 ("pandas.core.categorical", "Categorical"): ("pandas.core.arrays", "Categorical"), 

124 # 19939, add timedeltaindex, float64index compat from 15998 move 

125 ("pandas.tseries.tdi", "TimedeltaIndex"): ( 

126 "pandas.core.indexes.timedeltas", 

127 "TimedeltaIndex", 

128 ), 

129 ("pandas.indexes.numeric", "Float64Index"): ( 

130 "pandas.core.indexes.base", 

131 "Index", # updated in 50775 

132 ), 

133 # 50775, remove Int64Index, UInt64Index & Float64Index from codabase 

134 ("pandas.core.indexes.numeric", "Int64Index"): ( 

135 "pandas.core.indexes.base", 

136 "Index", 

137 ), 

138 ("pandas.core.indexes.numeric", "UInt64Index"): ( 

139 "pandas.core.indexes.base", 

140 "Index", 

141 ), 

142 ("pandas.core.indexes.numeric", "Float64Index"): ( 

143 "pandas.core.indexes.base", 

144 "Index", 

145 ), 

146 ("pandas.core.arrays.sparse.dtype", "SparseDtype"): ( 

147 "pandas.core.dtypes.dtypes", 

148 "SparseDtype", 

149 ), 

150} 

151 

152 

153# our Unpickler sub-class to override methods and some dispatcher 

154# functions for compat and uses a non-public class of the pickle module. 

155 

156 

157class Unpickler(pkl._Unpickler): 

158 def find_class(self, module, name): 

159 # override superclass 

160 key = (module, name) 

161 module, name = _class_locations_map.get(key, key) 

162 return super().find_class(module, name) 

163 

164 

165Unpickler.dispatch = copy.copy(Unpickler.dispatch) 

166Unpickler.dispatch[pkl.REDUCE[0]] = load_reduce 

167 

168 

169def load_newobj(self) -> None: 

170 args = self.stack.pop() 

171 cls = self.stack[-1] 

172 

173 # compat 

174 if issubclass(cls, Index): 

175 obj = object.__new__(cls) 

176 elif issubclass(cls, DatetimeArray) and not args: 

177 arr = np.array([], dtype="M8[ns]") 

178 obj = cls.__new__(cls, arr, arr.dtype) 

179 elif issubclass(cls, TimedeltaArray) and not args: 

180 arr = np.array([], dtype="m8[ns]") 

181 obj = cls.__new__(cls, arr, arr.dtype) 

182 elif cls is BlockManager and not args: 

183 obj = cls.__new__(cls, (), [], False) 

184 else: 

185 obj = cls.__new__(cls, *args) 

186 

187 self.stack[-1] = obj 

188 

189 

190Unpickler.dispatch[pkl.NEWOBJ[0]] = load_newobj 

191 

192 

193def load_newobj_ex(self) -> None: 

194 kwargs = self.stack.pop() 

195 args = self.stack.pop() 

196 cls = self.stack.pop() 

197 

198 # compat 

199 if issubclass(cls, Index): 

200 obj = object.__new__(cls) 

201 else: 

202 obj = cls.__new__(cls, *args, **kwargs) 

203 self.append(obj) 

204 

205 

206try: 

207 Unpickler.dispatch[pkl.NEWOBJ_EX[0]] = load_newobj_ex 

208except (AttributeError, KeyError): 

209 pass 

210 

211 

212def load(fh, encoding: str | None = None, is_verbose: bool = False): 

213 """ 

214 Load a pickle, with a provided encoding, 

215 

216 Parameters 

217 ---------- 

218 fh : a filelike object 

219 encoding : an optional encoding 

220 is_verbose : show exception output 

221 """ 

222 try: 

223 fh.seek(0) 

224 if encoding is not None: 

225 up = Unpickler(fh, encoding=encoding) 

226 else: 

227 up = Unpickler(fh) 

228 # "Unpickler" has no attribute "is_verbose" [attr-defined] 

229 up.is_verbose = is_verbose # type: ignore[attr-defined] 

230 

231 return up.load() 

232 except (ValueError, TypeError): 

233 raise 

234 

235 

236def loads( 

237 bytes_object: bytes, 

238 *, 

239 fix_imports: bool = True, 

240 encoding: str = "ASCII", 

241 errors: str = "strict", 

242): 

243 """ 

244 Analogous to pickle._loads. 

245 """ 

246 fd = io.BytesIO(bytes_object) 

247 return Unpickler( 

248 fd, fix_imports=fix_imports, encoding=encoding, errors=errors 

249 ).load() 

250 

251 

252@contextlib.contextmanager 

253def patch_pickle() -> Generator[None, None, None]: 

254 """ 

255 Temporarily patch pickle to use our unpickler. 

256 """ 

257 orig_loads = pkl.loads 

258 try: 

259 setattr(pkl, "loads", loads) 

260 yield 

261 finally: 

262 setattr(pkl, "loads", orig_loads)