pickle_compat.py 6.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229
  1. """
  2. Support pre-0.12 series pickle compatibility.
  3. """
  4. import copy
  5. import pickle as pkl
  6. import sys
  7. from pandas.compat import string_types, u # noqa
  8. import pandas # noqa
  9. from pandas import Index, compat
  10. def load_reduce(self):
  11. stack = self.stack
  12. args = stack.pop()
  13. func = stack[-1]
  14. if len(args) and type(args[0]) is type:
  15. n = args[0].__name__ # noqa
  16. try:
  17. stack[-1] = func(*args)
  18. return
  19. except Exception as e:
  20. # If we have a deprecated function,
  21. # try to replace and try again.
  22. msg = '_reconstruct: First argument must be a sub-type of ndarray'
  23. if msg in str(e):
  24. try:
  25. cls = args[0]
  26. stack[-1] = object.__new__(cls)
  27. return
  28. except TypeError:
  29. pass
  30. # try to re-encode the arguments
  31. if getattr(self, 'encoding', None) is not None:
  32. args = tuple(arg.encode(self.encoding)
  33. if isinstance(arg, string_types)
  34. else arg for arg in args)
  35. try:
  36. stack[-1] = func(*args)
  37. return
  38. except TypeError:
  39. pass
  40. # unknown exception, re-raise
  41. if getattr(self, 'is_verbose', None):
  42. print(sys.exc_info())
  43. print(func, args)
  44. raise
  45. # If classes are moved, provide compat here.
  46. _class_locations_map = {
  47. ('pandas.core.sparse.array', 'SparseArray'):
  48. ('pandas.core.arrays', 'SparseArray'),
  49. # 15477
  50. #
  51. # TODO: When FrozenNDArray is removed, add
  52. # the following lines for compat:
  53. #
  54. # ('pandas.core.base', 'FrozenNDArray'):
  55. # ('numpy', 'ndarray'),
  56. # ('pandas.core.indexes.frozen', 'FrozenNDArray'):
  57. # ('numpy', 'ndarray'),
  58. #
  59. # Afterwards, remove the current entry
  60. # for `pandas.core.base.FrozenNDArray`.
  61. ('pandas.core.base', 'FrozenNDArray'):
  62. ('pandas.core.indexes.frozen', 'FrozenNDArray'),
  63. ('pandas.core.base', 'FrozenList'):
  64. ('pandas.core.indexes.frozen', 'FrozenList'),
  65. # 10890
  66. ('pandas.core.series', 'TimeSeries'):
  67. ('pandas.core.series', 'Series'),
  68. ('pandas.sparse.series', 'SparseTimeSeries'):
  69. ('pandas.core.sparse.series', 'SparseSeries'),
  70. # 12588, extensions moving
  71. ('pandas._sparse', 'BlockIndex'):
  72. ('pandas._libs.sparse', 'BlockIndex'),
  73. ('pandas.tslib', 'Timestamp'):
  74. ('pandas._libs.tslib', 'Timestamp'),
  75. # 18543 moving period
  76. ('pandas._period', 'Period'): ('pandas._libs.tslibs.period', 'Period'),
  77. ('pandas._libs.period', 'Period'):
  78. ('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', '__nat_unpickle'),
  82. ('pandas._libs.tslib', '__nat_unpickle'):
  83. ('pandas._libs.tslibs.nattype', '__nat_unpickle'),
  84. # 15998 top-level dirs moving
  85. ('pandas.sparse.array', 'SparseArray'):
  86. ('pandas.core.arrays.sparse', 'SparseArray'),
  87. ('pandas.sparse.series', 'SparseSeries'):
  88. ('pandas.core.sparse.series', 'SparseSeries'),
  89. ('pandas.sparse.frame', 'SparseDataFrame'):
  90. ('pandas.core.sparse.frame', 'SparseDataFrame'),
  91. ('pandas.indexes.base', '_new_Index'):
  92. ('pandas.core.indexes.base', '_new_Index'),
  93. ('pandas.indexes.base', 'Index'):
  94. ('pandas.core.indexes.base', 'Index'),
  95. ('pandas.indexes.numeric', 'Int64Index'):
  96. ('pandas.core.indexes.numeric', 'Int64Index'),
  97. ('pandas.indexes.range', 'RangeIndex'):
  98. ('pandas.core.indexes.range', 'RangeIndex'),
  99. ('pandas.indexes.multi', 'MultiIndex'):
  100. ('pandas.core.indexes.multi', 'MultiIndex'),
  101. ('pandas.tseries.index', '_new_DatetimeIndex'):
  102. ('pandas.core.indexes.datetimes', '_new_DatetimeIndex'),
  103. ('pandas.tseries.index', 'DatetimeIndex'):
  104. ('pandas.core.indexes.datetimes', 'DatetimeIndex'),
  105. ('pandas.tseries.period', 'PeriodIndex'):
  106. ('pandas.core.indexes.period', 'PeriodIndex'),
  107. # 19269, arrays moving
  108. ('pandas.core.categorical', 'Categorical'):
  109. ('pandas.core.arrays', 'Categorical'),
  110. # 19939, add timedeltaindex, float64index compat from 15998 move
  111. ('pandas.tseries.tdi', 'TimedeltaIndex'):
  112. ('pandas.core.indexes.timedeltas', 'TimedeltaIndex'),
  113. ('pandas.indexes.numeric', 'Float64Index'):
  114. ('pandas.core.indexes.numeric', 'Float64Index'),
  115. }
  116. # our Unpickler sub-class to override methods and some dispatcher
  117. # functions for compat
  118. if compat.PY3:
  119. class Unpickler(pkl._Unpickler):
  120. def find_class(self, module, name):
  121. # override superclass
  122. key = (module, name)
  123. module, name = _class_locations_map.get(key, key)
  124. return super(Unpickler, self).find_class(module, name)
  125. else:
  126. class Unpickler(pkl.Unpickler):
  127. def find_class(self, module, name):
  128. # override superclass
  129. key = (module, name)
  130. module, name = _class_locations_map.get(key, key)
  131. __import__(module)
  132. mod = sys.modules[module]
  133. klass = getattr(mod, name)
  134. return klass
  135. Unpickler.dispatch = copy.copy(Unpickler.dispatch)
  136. Unpickler.dispatch[pkl.REDUCE[0]] = load_reduce
  137. def load_newobj(self):
  138. args = self.stack.pop()
  139. cls = self.stack[-1]
  140. # compat
  141. if issubclass(cls, Index):
  142. obj = object.__new__(cls)
  143. else:
  144. obj = cls.__new__(cls, *args)
  145. self.stack[-1] = obj
  146. Unpickler.dispatch[pkl.NEWOBJ[0]] = load_newobj
  147. def load_newobj_ex(self):
  148. kwargs = self.stack.pop()
  149. args = self.stack.pop()
  150. cls = self.stack.pop()
  151. # compat
  152. if issubclass(cls, Index):
  153. obj = object.__new__(cls)
  154. else:
  155. obj = cls.__new__(cls, *args, **kwargs)
  156. self.append(obj)
  157. try:
  158. Unpickler.dispatch[pkl.NEWOBJ_EX[0]] = load_newobj_ex
  159. except (AttributeError, KeyError):
  160. pass
  161. def load(fh, encoding=None, compat=False, is_verbose=False):
  162. """load a pickle, with a provided encoding
  163. if compat is True:
  164. fake the old class hierarchy
  165. if it works, then return the new type objects
  166. Parameters
  167. ----------
  168. fh : a filelike object
  169. encoding : an optional encoding
  170. compat : provide Series compatibility mode, boolean, default False
  171. is_verbose : show exception output
  172. """
  173. try:
  174. fh.seek(0)
  175. if encoding is not None:
  176. up = Unpickler(fh, encoding=encoding)
  177. else:
  178. up = Unpickler(fh)
  179. up.is_verbose = is_verbose
  180. return up.load()
  181. except (ValueError, TypeError):
  182. raise