Skip to content

Commit

Permalink
fix(io): use tmp dir fixture for compressedpickle doctests
Browse files Browse the repository at this point in the history
Signed-off-by: Cameron Smith <cameron.ray.smith@gmail.com>
  • Loading branch information
cameronraysmith committed Aug 2, 2024
1 parent 20953ae commit f678036
Showing 1 changed file with 14 additions and 9 deletions.
23 changes: 14 additions & 9 deletions src/pyrovelocity/io/compressedpickle.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import os
import pickle
from os import PathLike
from pathlib import Path
from typing import Any

import zstandard as zstd


__all__ = ["CompressedPickle"]


Expand All @@ -15,15 +14,17 @@ class CompressedPickle:
Examples:
>>> import pandas as pd
>>> tmp = getfixture("tmp_path")
>>> test_data_path = tmp / "test_data.pkl.zst"
>>> test_data = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
>>> CompressedPickle.save('test_data.pkl.zst', test_data)
>>> loaded_data = CompressedPickle.load('test_data.pkl.zst')
>>> CompressedPickle.save(test_data_path, test_data)
>>> loaded_data = CompressedPickle.load(test_data_path)
>>> loaded_data.equals(test_data)
True
"""

@staticmethod
def save(file_path: os.PathLike | str, obj: Any) -> None:
def save(file_path: PathLike | str, obj: Any) -> None:
"""
Save the given object to a zstandard-compressed pickle file.
Expand All @@ -33,8 +34,10 @@ def save(file_path: os.PathLike | str, obj: Any) -> None:
Examples:
>>> import pandas as pd
>>> tmp = getfixture("tmp_path")
>>> test_data_path = tmp / "test_data.pkl.zst"
>>> test_data = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
>>> CompressedPickle.save('test_data.pkl.zst', test_data)
>>> CompressedPickle.save(test_data_path, test_data)
"""
file_path = Path(file_path)
file_path.parent.mkdir(parents=True, exist_ok=True)
Expand All @@ -45,7 +48,7 @@ def save(file_path: os.PathLike | str, obj: Any) -> None:
pickle.dump(obj, compressor)

@staticmethod
def load(file_path: os.PathLike | str) -> Any:
def load(file_path: PathLike | str) -> Any:
"""
Load an object from a zstandard-compressed pickle file.
Expand All @@ -57,9 +60,11 @@ def load(file_path: os.PathLike | str) -> Any:
Examples:
>>> import pandas as pd
>>> tmp = getfixture("tmp_path")
>>> test_data_path = tmp / "test_data.pkl.zst"
>>> test_data = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
>>> CompressedPickle.save('test_data.pkl.zst', test_data)
>>> loaded_data = CompressedPickle.load('test_data.pkl.zst')
>>> CompressedPickle.save(test_data_path, test_data)
>>> loaded_data = CompressedPickle.load(test_data_path)
>>> loaded_data.equals(test_data)
True
"""
Expand Down

0 comments on commit f678036

Please sign in to comment.