143 lines
4.1 KiB
Python
143 lines
4.1 KiB
Python
"""
|
|
Pickle compatibility to pandas version 1.0
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import io
|
|
import pickle
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
)
|
|
|
|
import numpy as np
|
|
|
|
from pandas._libs.arrays import NDArrayBacked
|
|
from pandas._libs.tslibs import BaseOffset
|
|
|
|
from pandas.core.arrays import (
|
|
DatetimeArray,
|
|
PeriodArray,
|
|
TimedeltaArray,
|
|
)
|
|
from pandas.core.internals import BlockManager
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Generator
|
|
|
|
|
|
# If classes are moved, provide compat here.
|
|
_class_locations_map = {
|
|
# Re-routing unpickle block logic to go through _unpickle_block instead
|
|
# for pandas <= 1.3.5
|
|
("pandas.core.internals.blocks", "new_block"): (
|
|
"pandas._libs.internals",
|
|
"_unpickle_block",
|
|
),
|
|
# Avoid Cython's warning "contradiction to Python 'class private name' rules"
|
|
("pandas._libs.tslibs.nattype", "__nat_unpickle"): (
|
|
"pandas._libs.tslibs.nattype",
|
|
"_nat_unpickle",
|
|
),
|
|
# 50775, remove Int64Index, UInt64Index & Float64Index from codebase
|
|
("pandas.core.indexes.numeric", "Int64Index"): (
|
|
"pandas.core.indexes.base",
|
|
"Index",
|
|
),
|
|
("pandas.core.indexes.numeric", "UInt64Index"): (
|
|
"pandas.core.indexes.base",
|
|
"Index",
|
|
),
|
|
("pandas.core.indexes.numeric", "Float64Index"): (
|
|
"pandas.core.indexes.base",
|
|
"Index",
|
|
),
|
|
("pandas.core.arrays.sparse.dtype", "SparseDtype"): (
|
|
"pandas.core.dtypes.dtypes",
|
|
"SparseDtype",
|
|
),
|
|
}
|
|
|
|
|
|
# our Unpickler sub-class to override methods and some dispatcher
|
|
# functions for compat and uses a non-public class of the pickle module.
|
|
class Unpickler(pickle._Unpickler):
|
|
def find_class(self, module: str, name: str) -> Any:
|
|
key = (module, name)
|
|
module, name = _class_locations_map.get(key, key)
|
|
return super().find_class(module, name)
|
|
|
|
dispatch = pickle._Unpickler.dispatch.copy()
|
|
|
|
def load_reduce(self) -> None:
|
|
stack = self.stack # type: ignore[attr-defined]
|
|
args = stack.pop()
|
|
func = stack[-1]
|
|
|
|
try:
|
|
stack[-1] = func(*args)
|
|
except TypeError:
|
|
# If we have a deprecated function,
|
|
# try to replace and try again.
|
|
if args and isinstance(args[0], type) and issubclass(args[0], BaseOffset):
|
|
# TypeError: object.__new__(Day) is not safe, use Day.__new__()
|
|
cls = args[0]
|
|
stack[-1] = cls.__new__(*args)
|
|
return
|
|
elif args and issubclass(args[0], PeriodArray):
|
|
cls = args[0]
|
|
stack[-1] = NDArrayBacked.__new__(*args)
|
|
return
|
|
raise
|
|
|
|
dispatch[pickle.REDUCE[0]] = load_reduce # type: ignore[assignment]
|
|
|
|
def load_newobj(self) -> None:
|
|
args = self.stack.pop() # type: ignore[attr-defined]
|
|
cls = self.stack.pop() # type: ignore[attr-defined]
|
|
|
|
# compat
|
|
if issubclass(cls, DatetimeArray) and not args:
|
|
arr = np.array([], dtype="M8[ns]")
|
|
obj = cls.__new__(cls, arr, arr.dtype)
|
|
elif issubclass(cls, TimedeltaArray) and not args:
|
|
arr = np.array([], dtype="m8[ns]")
|
|
obj = cls.__new__(cls, arr, arr.dtype)
|
|
elif cls is BlockManager and not args:
|
|
obj = cls.__new__(cls, (), [], False)
|
|
else:
|
|
obj = cls.__new__(cls, *args)
|
|
self.append(obj) # type: ignore[attr-defined]
|
|
|
|
dispatch[pickle.NEWOBJ[0]] = load_newobj # type: ignore[assignment]
|
|
|
|
|
|
def loads(
|
|
bytes_object: bytes,
|
|
*,
|
|
fix_imports: bool = True,
|
|
encoding: str = "ASCII",
|
|
errors: str = "strict",
|
|
) -> Any:
|
|
"""
|
|
Analogous to pickle._loads.
|
|
"""
|
|
fd = io.BytesIO(bytes_object)
|
|
return Unpickler(
|
|
fd, fix_imports=fix_imports, encoding=encoding, errors=errors
|
|
).load()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def patch_pickle() -> Generator[None]:
|
|
"""
|
|
Temporarily patch pickle to use our unpickler.
|
|
"""
|
|
orig_loads = pickle.loads
|
|
try:
|
|
setattr(pickle, "loads", loads)
|
|
yield
|
|
finally:
|
|
setattr(pickle, "loads", orig_loads)
|