Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-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 Generator 

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 

25 

26def load_reduce(self): 

27 stack = self.stack 

28 args = stack.pop() 

29 func = stack[-1] 

30 

31 try: 

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

33 return 

34 except TypeError as err: 

35 # If we have a deprecated function, 

36 # try to replace and try again. 

37 

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

39 

40 if msg in str(err): 

41 try: 

42 cls = args[0] 

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

44 return 

45 except TypeError: 

46 pass 

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

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

49 cls = args[0] 

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

51 return 

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

53 cls = args[0] 

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

55 return 

56 

57 raise 

58 

59 

60# If classes are moved, provide compat here. 

61_class_locations_map = { 

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

63 # 15477 

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

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

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

67 # 10890 

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

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

70 "pandas.core.sparse.series", 

71 "SparseSeries", 

72 ), 

73 # 12588, extensions moving 

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

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

76 # 18543 moving period 

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

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

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

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

81 "pandas._libs.tslibs.nattype", 

82 "__nat_unpickle", 

83 ), 

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

85 "pandas._libs.tslibs.nattype", 

86 "__nat_unpickle", 

87 ), 

88 # 15998 top-level dirs moving 

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

90 "pandas.core.arrays.sparse", 

91 "SparseArray", 

92 ), 

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

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

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

96 "pandas.core.indexes.base", 

97 "Index", # updated in 50775 

98 ), 

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

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

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

102 "pandas.core.indexes.datetimes", 

103 "_new_DatetimeIndex", 

104 ), 

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

106 "pandas.core.indexes.datetimes", 

107 "DatetimeIndex", 

108 ), 

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

110 "pandas.core.indexes.period", 

111 "PeriodIndex", 

112 ), 

113 # 19269, arrays moving 

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

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

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

117 "pandas.core.indexes.timedeltas", 

118 "TimedeltaIndex", 

119 ), 

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

121 "pandas.core.indexes.base", 

122 "Index", # updated in 50775 

123 ), 

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

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

126 "pandas.core.indexes.base", 

127 "Index", 

128 ), 

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

130 "pandas.core.indexes.base", 

131 "Index", 

132 ), 

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

134 "pandas.core.indexes.base", 

135 "Index", 

136 ), 

137} 

138 

139 

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

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

142 

143 

144class Unpickler(pkl._Unpickler): 

145 def find_class(self, module, name): 

146 # override superclass 

147 key = (module, name) 

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

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

150 

151 

152Unpickler.dispatch = copy.copy(Unpickler.dispatch) 

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

154 

155 

156def load_newobj(self) -> None: 

157 args = self.stack.pop() 

158 cls = self.stack[-1] 

159 

160 # compat 

161 if issubclass(cls, Index): 

162 obj = object.__new__(cls) 

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

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

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

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

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

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

169 elif cls is BlockManager and not args: 

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

171 else: 

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

173 

174 self.stack[-1] = obj 

175 

176 

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

178 

179 

180def load_newobj_ex(self) -> None: 

181 kwargs = self.stack.pop() 

182 args = self.stack.pop() 

183 cls = self.stack.pop() 

184 

185 # compat 

186 if issubclass(cls, Index): 

187 obj = object.__new__(cls) 

188 else: 

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

190 self.append(obj) 

191 

192 

193try: 

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

195except (AttributeError, KeyError): 

196 pass 

197 

198 

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

200 """ 

201 Load a pickle, with a provided encoding, 

202 

203 Parameters 

204 ---------- 

205 fh : a filelike object 

206 encoding : an optional encoding 

207 is_verbose : show exception output 

208 """ 

209 try: 

210 fh.seek(0) 

211 if encoding is not None: 

212 up = Unpickler(fh, encoding=encoding) 

213 else: 

214 up = Unpickler(fh) 

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

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

217 

218 return up.load() 

219 except (ValueError, TypeError): 

220 raise 

221 

222 

223def loads( 

224 bytes_object: bytes, 

225 *, 

226 fix_imports: bool = True, 

227 encoding: str = "ASCII", 

228 errors: str = "strict", 

229): 

230 """ 

231 Analogous to pickle._loads. 

232 """ 

233 fd = io.BytesIO(bytes_object) 

234 return Unpickler( 

235 fd, fix_imports=fix_imports, encoding=encoding, errors=errors 

236 ).load() 

237 

238 

239@contextlib.contextmanager 

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

241 """ 

242 Temporarily patch pickle to use our unpickler. 

243 """ 

244 orig_loads = pkl.loads 

245 try: 

246 setattr(pkl, "loads", loads) 

247 yield 

248 finally: 

249 setattr(pkl, "loads", orig_loads)