"""数据IO工具.
提供 pickle 文件的读写功能，支持多种序列化引擎和压缩格式。
"""
import gzip
import pickle
from contextlib import nullcontext
from io import BytesIO
from pathlib import Path
from typing import Any, Optional, Union
import joblib
from ..exceptions import DependencyError, ValidationError
def _open_file(
file: Union[str, Path],
mode: str = 'rb',
compression: Optional[str] = None
):
"""根据压缩格式打开文件."""
if compression == 'gzip' or str(file).lower().endswith('.gz'):
return gzip.open(file, mode)
elif compression == 'bz2' or str(file).lower().endswith('.bz2'):
import bz2
return bz2.open(file, mode)
elif compression == 'xz' or str(file).lower().endswith('.xz'):
try:
import lzma
return lzma.open(file, mode)
except ImportError:
raise DependencyError("缺少可选依赖 lzma，请改用 gzip 或 bz2 压缩格式")
elif compression == 'lz4' or str(file).lower().endswith('.lz4'):
try:
import lz4.frame
return lz4.frame.open(file, mode)
except ImportError:
raise DependencyError("缺少可选依赖 lz4，请先安装: pip install lz4")
elif compression in ('zstd', 'zstandard') or str(file).lower().endswith(('.zst', '.zstd')):
try:
import zstandard
return zstandard.open(file, mode)
except ImportError:
raise DependencyError("缺少可选依赖 zstandard，请先安装: pip install zstandard")
else:
return nullcontext(file) if hasattr(file, 'read') else open(file, mode)
[文档]
def load_pickle(
file: Union[str, Path],
engine: str = "auto",
compression: Optional[str] = None,
*,
source_name: Optional[str] = None,
) -> Any:
"""导入 pickle 文件。
支持多种序列化引擎（joblib/dill/cloudpickle/pickle）和压缩格式
（gzip/bz2/xz/lz4/zstd）。支持根据文件扩展名自动检测。
:param file: pickle 文件路径或可定位的二进制文件对象
:param source_name: 文件对象的原始文件名，仅用于格式推断，不会重新打开该路径
:param engine: 使用的序列化引擎，可选：
- 'auto': 自动检测（根据文件内容和扩展名推断，默认）
- 'joblib': 使用 joblib（推荐用于 numpy/scipy/sklearn 对象）
- 'dill': 使用 dill（支持 lambda、嵌套函数等复杂对象）
- 'cloudpickle': 使用 cloudpickle（常用于分布式计算如 PyTorch/Spark）
- 'pickle': 使用标准库 pickle
:param compression: 压缩格式，可选：
- None: 根据文件扩展名自动检测（.gz/.bz2/.xz/.lz4/.zst）
- 'gzip'/'gz': gzip 压缩
- 'bz2': bzip2 压缩
- 'xz': xz/lzma 压缩
- 'lz4': lz4 压缩（需安装 lz4）
- 'zstd'/'zstandard': zstd 压缩（需安装 zstandard）
:return: 反序列化后的对象
**参考样例**
>>> data = load_pickle('model.pkl')
>>> data = load_pickle('model.pkl.gz')
>>> data = load_pickle('model.dill', engine='dill')
>>> data = load_pickle('model.pkl', engine='cloudpickle')
>>> data = load_pickle('model.pkl', compression='gzip')
"""
file_str = str(source_name if source_name is not None else file).lower()
# 自动检测压缩格式
comp = compression
if comp is None:
if file_str.endswith('.gz') or file_str.endswith('.gzip'):
comp = 'gzip'
elif file_str.endswith('.bz2'):
comp = 'bz2'
elif file_str.endswith('.xz'):
comp = 'xz'
elif file_str.endswith('.lz4'):
comp = 'lz4'
elif file_str.endswith('.zst') or file_str.endswith('.zstd'):
comp = 'zstd'
# 自动检测引擎
eng = engine
if eng == "auto":
# 根据文件扩展名推断
if file_str.endswith('.joblib') or file_str.endswith('.joblib.gz'):
eng = 'joblib'
elif file_str.endswith('.dill') or file_str.endswith('.dill.gz'):
eng = 'dill'
elif file_str.endswith('.cloudpickle'):
eng = 'cloudpickle'
else:
# 默认使用 joblib（业内最常用）
eng = 'joblib'
# 使用指定引擎加载
if eng == "joblib":
if comp:
# 直接读取解压流，避免再分配一份完整解压字节缓冲。
with _open_file(file, 'rb', comp) as f:
return joblib.load(f)
return joblib.load(file)
elif eng == "dill":
try:
import dill
with _open_file(file, "rb", comp) as f:
return dill.load(f)
except ImportError:
raise DependencyError("缺少可选依赖 dill，请先安装: pip install dill")
elif eng == "cloudpickle":
try:
import cloudpickle
with _open_file(file, "rb", comp) as f:
return cloudpickle.load(f)
except ImportError:
raise DependencyError("缺少可选依赖 cloudpickle，请先安装: pip install cloudpickle")
elif eng == "pickle":
with _open_file(file, "rb", comp) as f:
return pickle.load(f)
else:
raise ValidationError(
f"engine 目前只支持 ['auto', 'joblib', 'dill', 'cloudpickle', 'pickle'], "
f"不支持 {eng}"
)
[文档]
def save_pickle(
obj: Any,
file: Union[str, Path],
engine: str = "joblib",
compression: Optional[str] = None,
compression_level: Optional[int] = None,
protocol: Optional[int] = None
) -> str:
"""保存数据至 pickle 文件。
支持多种序列化引擎（joblib/dill/cloudpickle/pickle）和压缩格式
（gzip/bz2/xz/lz4/zstd），可处理大型模型和复杂对象。
:param obj: 需要保存的数据对象
:param file: 文件路径，建议扩展名 .pkl, .joblib, .dill 等
:param engine: 使用的序列化引擎，可选：
- 'joblib': joblib（默认，推荐用于 numpy/scipy/sklearn 对象）
- 'dill': dill（支持 lambda、嵌套函数等复杂对象）
- 'cloudpickle': cloudpickle（常用于分布式计算）
- 'pickle': 标准库 pickle
:param compression: 压缩格式，可选：
- None: 不压缩（默认）
- 'gzip'/'gz': gzip 压缩（兼容性好）
- 'bz2': bzip2 压缩（压缩率高但较慢）
- 'xz': xz/lzma 压缩（最高压缩率）
- 'lz4': lz4 压缩（速度最快，需安装 lz4）
- 'zstd'/'zstandard': zstd 压缩（速度与压缩率平衡，需安装 zstandard）
- 'auto': 根据文件扩展名自动选择
:param compression_level: 压缩级别（1-9，数字越大压缩率越高，默认取决于压缩算法）
:param protocol: pickle 协议版本（默认使用最高可用版本）
:return: 保存的文件路径
**参考样例**
>>> save_pickle(model, 'model.pkl')
>>> save_pickle(lambda_func, 'func.dill', engine='dill')
>>> save_pickle(model, 'model.pkl', engine='cloudpickle')
>>> save_pickle(model, 'model.pkl.gz')
>>> save_pickle(model, 'model.pkl', compression='zstd', compression_level=3)
>>> save_pickle(model, 'model.pkl.xz', compression='xz')
"""
file_str = str(file).lower()
# 自动检测压缩格式
comp = compression
if comp == 'auto' or comp is None:
if file_str.endswith('.gz') or file_str.endswith('.gzip'):
comp = 'gzip'
elif file_str.endswith('.bz2'):
comp = 'bz2'
elif file_str.endswith('.xz'):
comp = 'xz'
elif file_str.endswith('.lz4'):
comp = 'lz4'
elif file_str.endswith('.zst') or file_str.endswith('.zstd'):
comp = 'zstd'
# 设置 pickle protocol
proto = protocol
if proto is None:
proto = pickle.HIGHEST_PROTOCOL
# 序列化
if engine == "joblib":
if comp:
# joblib 对压缩支持有限，先序列化到内存再压缩
buf = BytesIO()
joblib.dump(obj, buf)
buf.seek(0)
with _open_file(file, "wb", comp) as f:
f.write(buf.getvalue()) # type: ignore
else:
joblib.dump(obj, file)
elif engine == "dill":
try:
import dill
with _open_file(file, "wb", comp) as f:
dill.dump(obj, f, protocol=proto)
except ImportError:
raise DependencyError("缺少可选依赖 dill，请先安装: pip install dill")
elif engine == "cloudpickle":
try:
import cloudpickle
with _open_file(file, "wb", comp) as f:
cloudpickle.dump(obj, f, protocol=proto)
except ImportError:
raise DependencyError("缺少可选依赖 cloudpickle，请先安装: pip install cloudpickle")
elif engine == "pickle":
with _open_file(file, "wb", comp) as f:
pickle.dump(obj, f, protocol=proto)
else:
raise ValidationError(
f"engine 目前只支持 ['joblib', 'dill', 'cloudpickle', 'pickle'], "
f"不支持 {engine}"
)
return str(file)