Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Pickle compatibility to pandas version 1.0
3"""
5from __future__ import annotations
7import contextlib
8import io
9import pickle
10from typing import (
11 TYPE_CHECKING,
12 Any,
13)
15import numpy as np
17from pandas._libs.arrays import NDArrayBacked
18from pandas._libs.tslibs import BaseOffset
20from pandas.core.arrays import (
21 DatetimeArray,
22 PeriodArray,
23 TimedeltaArray,
24)
25from pandas.core.internals import BlockManager
27if TYPE_CHECKING:
28 from collections.abc import Generator
31# If classes are moved, provide compat here.
32_class_locations_map = {
33 # Re-routing unpickle block logic to go through _unpickle_block instead
34 # for pandas <= 1.3.5
35 ("pandas.core.internals.blocks", "new_block"): (
36 "pandas._libs.internals",
37 "_unpickle_block",
38 ),
39 # Avoid Cython's warning "contradiction to Python 'class private name' rules"
40 ("pandas._libs.tslibs.nattype", "__nat_unpickle"): (
41 "pandas._libs.tslibs.nattype",
42 "_nat_unpickle",
43 ),
44 # 50775, remove Int64Index, UInt64Index & Float64Index from codebase
45 ("pandas.core.indexes.numeric", "Int64Index"): (
46 "pandas.core.indexes.base",
47 "Index",
48 ),
49 ("pandas.core.indexes.numeric", "UInt64Index"): (
50 "pandas.core.indexes.base",
51 "Index",
52 ),
53 ("pandas.core.indexes.numeric", "Float64Index"): (
54 "pandas.core.indexes.base",
55 "Index",
56 ),
57 ("pandas.core.arrays.sparse.dtype", "SparseDtype"): (
58 "pandas.core.dtypes.dtypes",
59 "SparseDtype",
60 ),
61}
64# our Unpickler sub-class to override methods and some dispatcher
65# functions for compat and uses a non-public class of the pickle module.
66class Unpickler(pickle._Unpickler):
67 def find_class(self, module: str, name: str) -> Any:
68 key = (module, name)
69 module, name = _class_locations_map.get(key, key)
70 return super().find_class(module, name)
72 dispatch = pickle._Unpickler.dispatch.copy()
74 def load_reduce(self) -> None:
75 stack = self.stack # type: ignore[attr-defined]
76 args = stack.pop()
77 func = stack[-1]
79 try:
80 stack[-1] = func(*args)
81 except TypeError:
82 # If we have a deprecated function,
83 # try to replace and try again.
84 if args and isinstance(args[0], type) and issubclass(args[0], BaseOffset):
85 # TypeError: object.__new__(Day) is not safe, use Day.__new__()
86 cls = args[0]
87 stack[-1] = cls.__new__(*args)
88 return
89 elif args and issubclass(args[0], PeriodArray):
90 cls = args[0]
91 stack[-1] = NDArrayBacked.__new__(*args)
92 return
93 raise
95 dispatch[pickle.REDUCE[0]] = load_reduce # type: ignore[assignment]
97 def load_newobj(self) -> None:
98 args = self.stack.pop() # type: ignore[attr-defined]
99 cls = self.stack.pop() # type: ignore[attr-defined]
101 # compat
102 if issubclass(cls, DatetimeArray) and not args:
103 arr = np.array([], dtype="M8[ns]")
104 obj = cls.__new__(cls, arr, arr.dtype)
105 elif issubclass(cls, TimedeltaArray) and not args:
106 arr = np.array([], dtype="m8[ns]")
107 obj = cls.__new__(cls, arr, arr.dtype)
108 elif cls is BlockManager and not args:
109 obj = cls.__new__(cls, (), [], False)
110 else:
111 obj = cls.__new__(cls, *args)
112 self.append(obj) # type: ignore[attr-defined]
114 dispatch[pickle.NEWOBJ[0]] = load_newobj # type: ignore[assignment]
117def loads(
118 bytes_object: bytes,
119 *,
120 fix_imports: bool = True,
121 encoding: str = "ASCII",
122 errors: str = "strict",
123) -> Any:
124 """
125 Analogous to pickle._loads.
126 """
127 fd = io.BytesIO(bytes_object)
128 return Unpickler(
129 fd, fix_imports=fix_imports, encoding=encoding, errors=errors
130 ).load()
133@contextlib.contextmanager
134def patch_pickle() -> Generator[None]:
135 """
136 Temporarily patch pickle to use our unpickler.
137 """
138 orig_loads = pickle.loads
139 try:
140 setattr(pickle, "loads", loads)
141 yield
142 finally:
143 setattr(pickle, "loads", orig_loads)