Visitar URL original
Refactor pickle loading to handle paths from different operating syst… · DeepLabCut/DeepLabCut@594ba9f · GitHub
Skip to content

Commit 594ba9f

Browse files
authored
Refactor pickle loading to handle paths from different operating systems (#3565)
1 parent beb39f6 commit 594ba9f

3 files changed

Lines changed: 46 additions & 6 deletions

File tree

‎deeplabcut/utils/auxfun_multianimal.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -211,8 +211,7 @@ def LoadFullMultiAnimalData(dataname):
211211
data = pickle.load(handle)
212212
except (pickle.UnpicklingError, FileNotFoundError):
213213
data = shelve.open(data_file, flag="r")
214-
with metadata_file.open("rb") as handle:
215-
metadata = pickle.load(handle)
214+
metadata = auxiliaryfunctions.read_pickle(metadata_file)
216215
return data, metadata
217216

218217

‎deeplabcut/utils/auxiliaryfunctions.py‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import os
2525
import pickle
2626
from collections.abc import Sequence
27-
from pathlib import Path
27+
from pathlib import Path, PurePosixPath, PureWindowsPath
2828

2929
import pandas as pd
3030

@@ -139,10 +139,28 @@ def attempt_to_make_folder(foldername, recursive=False):
139139
foldername.mkdir(exist_ok=True, parents=True)
140140

141141

142+
_FOREIGN_PATH_NAME, _FOREIGN_PATH_TYPE = (
143+
("PosixPath", PurePosixPath) if os.name == "nt" else ("WindowsPath", PureWindowsPath)
144+
)
145+
146+
147+
class _PortableUnpickler(pickle.Unpickler):
148+
"""Unpickler loading paths pickled on another OS as pure paths."""
149+
150+
def find_class(self, module, name):
151+
# Python 3.13+ pickles paths from ``pathlib._local``
152+
if module in ("pathlib", "pathlib._local") and name == _FOREIGN_PATH_NAME:
153+
return _FOREIGN_PATH_TYPE
154+
return super().find_class(module, name)
155+
156+
142157
def read_pickle(filename):
143-
"""Read the pickle file."""
158+
"""Read the pickle file.
159+
160+
Paths pickled on another OS are loaded as ``PurePosixPath`` or ``PureWindowsPath``.
161+
"""
144162
with Path(filename).open("rb") as handle:
145-
return pickle.load(handle)
163+
return _PortableUnpickler(handle).load()
146164

147165

148166
def write_pickle(filename, data):

‎tests/test_auxiliaryfunctions.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,8 @@
88
#
99
# Licensed under GNU Lesser General Public License v3.0
1010
#
11-
from pathlib import Path
11+
import os
12+
from pathlib import Path, PurePosixPath, PureWindowsPath
1213
from unittest.mock import patch
1314

1415
import pytest
@@ -349,3 +350,25 @@ def test_project_path_updated_when_moved(self, tmp_path):
349350

350351
expected = project_dir.absolute()
351352
assert cfg["project_path"] == expected
353+
354+
355+
def _write_path_pickle(path: Path, module: str, cls_name: str) -> Path:
356+
# Protocol 0 pickle of ``{"project_path": <module>.<cls_name>("a", "b")}``; written by hand
357+
# because a concrete path of the other OS cannot be created to pickle it
358+
path.write_bytes(f"(dVproject_path\nc{module}\n{cls_name}\n(Va\nVb\ntRs.".encode())
359+
return path
360+
361+
362+
@pytest.mark.parametrize("module", ["pathlib", "pathlib._local"])
363+
def test_read_pickle_loads_foreign_paths_as_pure(tmp_path, module):
364+
"""Paths pickled on another OS load as pure paths instead of raising (#3562)."""
365+
foreign_name, pure_type = ("PosixPath", PurePosixPath) if os.name == "nt" else ("WindowsPath", PureWindowsPath)
366+
data = auxiliaryfunctions.read_pickle(_write_path_pickle(tmp_path / "meta.pickle", module, foreign_name))
367+
assert type(data["project_path"]) is pure_type
368+
assert data["project_path"].parts == ("a", "b")
369+
370+
371+
def test_read_pickle_keeps_native_paths(tmp_path):
372+
native_type = type(Path())
373+
data = auxiliaryfunctions.read_pickle(_write_path_pickle(tmp_path / "meta.pickle", "pathlib", native_type.__name__))
374+
assert type(data["project_path"]) is native_type

0 commit comments

Comments
 (0)