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 datetime
9import io
10import pickle
11from typing import (
12 TYPE_CHECKING,
13 Any,
14)
16import numpy as np
18from pandas._libs.arrays import NDArrayBacked
19from pandas._libs.tslibs import (
20 BaseOffset,
21 Timestamp,
22)
24from pandas.core.arrays import (
25 DatetimeArray,
26 PeriodArray,
27 TimedeltaArray,
28)
29from pandas.core.internals import BlockManager
31if TYPE_CHECKING:
32 from collections.abc import Generator
35# If classes are moved, provide compat here.
36_class_locations_map = {
37 # Re-routing unpickle block logic to go through _unpickle_block instead
38 # for pandas <= 1.3.5
39 ("pandas.core.internals.blocks", "new_block"): (
40 "pandas._libs.internals",
41 "_unpickle_block",
42 ),
43 # Avoid Cython's warning "contradiction to Python 'class private name' rules"
44 ("pandas._libs.tslibs.nattype", "__nat_unpickle"): (
45 "pandas._libs.tslibs.nattype",
46 "_nat_unpickle",
47 ),
48 # 50775, remove Int64Index, UInt64Index & Float64Index from codebase
49 ("pandas.core.indexes.numeric", "Int64Index"): (
50 "pandas.core.indexes.base",
51 "Index",
52 ),
53 ("pandas.core.indexes.numeric", "UInt64Index"): (
54 "pandas.core.indexes.base",
55 "Index",
56 ),
57 ("pandas.core.indexes.numeric", "Float64Index"): (
58 "pandas.core.indexes.base",
59 "Index",
60 ),
61 ("pandas.core.arrays.sparse.dtype", "SparseDtype"): (
62 "pandas.core.dtypes.dtypes",
63 "SparseDtype",
64 ),
65}
68# our Unpickler sub-class to override methods and some dispatcher
69# functions for compat and uses a non-public class of the pickle module.
70class Unpickler(pickle._Unpickler):
71 def find_class(self, module: str, name: str) -> Any:
72 key = (module, name)
73 module, name = _class_locations_map.get(key, key)
74 return super().find_class(module, name)
76 dispatch = pickle._Unpickler.dispatch.copy()
78 def load_reduce(self) -> None:
79 stack = self.stack # type: ignore[attr-defined]
80 args = stack.pop()
81 func = stack[-1]
83 if (
84 func is Timestamp
85 and len(args) == 3
86 and (args[1] is None or isinstance(args[1], (str, BaseOffset)))
87 and (args[2] is None or isinstance(args[2], datetime.tzinfo))
88 ):
89 # GH#61792 pandas <= 1.2 pickled Timestamp as (value, freq, tz).
90 # `freq` is gone, so those slots now mean year/month and passing
91 # args positionally silently drops the tz. The isinstance checks
92 # keep a genuine by-component reduce, e.g. (Timestamp, (2020, 1, 2)),
93 # out of this branch.
94 value, _freq, tz = args
95 # An explicit tz=None rejects a tz-aware value, so only pass it when set.
96 stack[-1] = Timestamp(value) if tz is None else Timestamp(value, tz=tz)
97 return
99 try:
100 stack[-1] = func(*args)
101 except TypeError:
102 # If we have a deprecated function,
103 # try to replace and try again.
104 if args and isinstance(args[0], type) and issubclass(args[0], BaseOffset):
105 # TypeError: object.__new__(Day) is not safe, use Day.__new__()
106 cls = args[0]
107 stack[-1] = cls.__new__(*args)
108 return
109 elif args and issubclass(args[0], PeriodArray):
110 cls = args[0]
111 stack[-1] = NDArrayBacked.__new__(*args)
112 return
113 raise
115 dispatch[pickle.REDUCE[0]] = load_reduce # type: ignore[assignment]
117 def load_newobj(self) -> None:
118 args = self.stack.pop() # type: ignore[attr-defined]
119 cls = self.stack.pop() # type: ignore[attr-defined]
121 # compat
122 if issubclass(cls, DatetimeArray) and not args:
123 arr = np.array([], dtype="M8[ns]")
124 obj = cls.__new__(cls, arr, arr.dtype) # pyright: ignore[reportCallIssue]
125 elif issubclass(cls, TimedeltaArray) and not args:
126 arr = np.array([], dtype="m8[ns]")
127 obj = cls.__new__(cls, arr, arr.dtype) # pyright: ignore[reportCallIssue]
128 elif cls is BlockManager and not args:
129 # error: Too many arguments for "__new__" of "object"
130 obj = cls.__new__(cls, (), [], False) # type: ignore[call-arg]
131 else:
132 obj = cls.__new__(cls, *args)
133 self.append(obj) # type: ignore[attr-defined]
135 dispatch[pickle.NEWOBJ[0]] = load_newobj # type: ignore[assignment]
138def loads(
139 bytes_object: bytes,
140 *,
141 fix_imports: bool = True,
142 encoding: str = "ASCII",
143 errors: str = "strict",
144) -> Any:
145 """
146 Analogous to pickle._loads.
147 """
148 fd = io.BytesIO(bytes_object)
149 return Unpickler(
150 fd, fix_imports=fix_imports, encoding=encoding, errors=errors
151 ).load()
154@contextlib.contextmanager
155def patch_pickle() -> Generator[None]:
156 """
157 Temporarily patch pickle to use our unpickler.
158 """
159 orig_loads = pickle.loads
160 try:
161 setattr(pickle, "loads", loads)
162 yield
163 finally:
164 setattr(pickle, "loads", orig_loads)