diff --git a/pandas-stubs/_typing.pyi b/pandas-stubs/_typing.pyi index ac175a4b1..0a3473a40 100644 --- a/pandas-stubs/_typing.pyi +++ b/pandas-stubs/_typing.pyi @@ -67,6 +67,9 @@ class BaseBuffer(Protocol): ... class ReadBuffer(BaseBuffer, Protocol[AnyStr_cov]): ... class WriteBuffer(BaseBuffer, Protocol[AnyStr_cov]): ... +class ReadPickleBuffer(ReadBuffer[bytes], Protocol): + def readline(self, size: int | None = ...) -> bytes: ... + class ReadCsvBuffer(ReadBuffer[AnyStr_cov], Protocol[AnyStr_cov]): def __iter__(self) -> Iterator[AnyStr_cov]: ... def readline(self) -> AnyStr_cov: ... diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index a0e46ba9e..3391c4aaf 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -1947,12 +1947,6 @@ class DataFrame(NDFrame, OpsMixin): index: _bool = ..., indent: int | None = ..., ) -> _str: ... - def to_pickle( - self, - path: _str, - compression: CompressionOptions = ..., - protocol: int = ..., - ) -> None: ... @overload def to_string( self, diff --git a/pandas-stubs/core/generic.pyi b/pandas-stubs/core/generic.pyi index 56a78a703..c072be509 100644 --- a/pandas-stubs/core/generic.pyi +++ b/pandas-stubs/core/generic.pyi @@ -42,6 +42,7 @@ from pandas._typing import ( SortKind, StorageOptions, T, + WriteBuffer, ) from pandas.io.pytables import HDFStore @@ -171,9 +172,10 @@ class NDFrame(PandasObject, indexing.IndexingMixin): ) -> int | None: ... def to_pickle( self, - path: _str, + path: FilePath | WriteBuffer[bytes], compression: CompressionOptions = ..., protocol: int = ..., + storage_options: StorageOptions = ..., ) -> None: ... def to_clipboard( self, excel: _bool = ..., sep: _str | None = ..., **kwargs diff --git a/pandas-stubs/io/pickle.pyi b/pandas-stubs/io/pickle.pyi index 1f7774fe6..b2ead7903 100644 --- a/pandas-stubs/io/pickle.pyi +++ b/pandas-stubs/io/pickle.pyi @@ -1,15 +1,22 @@ +from typing import Any + from pandas._typing import ( CompressionOptions, - FilePathOrBuffer, + FilePath, + ReadPickleBuffer, + StorageOptions, + WriteBuffer, ) def to_pickle( - obj, - filepath_or_buffer: FilePathOrBuffer, - compression: str | None = ..., + obj: object, + filepath_or_buffer: FilePath | WriteBuffer[bytes], + compression: CompressionOptions = ..., protocol: int = ..., -): ... + storage_options: StorageOptions = ..., +) -> None: ... def read_pickle( - filepath_or_buffer_or_reader: FilePathOrBuffer, + filepath_or_buffer: FilePath | ReadPickleBuffer, compression: CompressionOptions = ..., -): ... + storage_options: StorageOptions = ..., +) -> Any: ... diff --git a/tests/test_io.py b/tests/test_io.py index 478add136..0e0f47b47 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -6,6 +6,7 @@ import sqlite3 from typing import ( TYPE_CHECKING, + Any, Dict, Generator, List, @@ -30,6 +31,7 @@ read_json, read_orc, read_parquet, + read_pickle, read_sas, read_spss, read_sql, @@ -45,6 +47,7 @@ from tests import check +from pandas.io.api import to_pickle from pandas.io.clipboard import PyperclipException from pandas.io.common import IOHandles from pandas.io.json._json import JsonReader @@ -117,6 +120,67 @@ def test_xml_str(): check(assert_type(read_xml(io.StringIO(out)), DataFrame), DataFrame) +def test_pickle(): + with ensure_clean() as path: + check(assert_type(DF.to_pickle(path), None), type(None)) + check(assert_type(read_pickle(path), Any), DataFrame) + + with ensure_clean() as path: + check(assert_type(to_pickle(DF, path), None), type(None)) + check(assert_type(read_pickle(path), Any), DataFrame) + + +def test_pickle_file_handle(): + with ensure_clean() as path: + check(assert_type(DF.to_pickle(path), None), type(None)) + file = open(path, "rb") + check(assert_type(read_pickle(file), Any), DataFrame) + file.close() + + +def test_pickle_path(): + with ensure_clean() as path: + check(assert_type(DF.to_pickle(path), None), type(None)) + check(assert_type(read_pickle(Path(path)), Any), DataFrame) + + +def test_pickle_protocol(): + with ensure_clean() as path: + DF.to_pickle(path, protocol=3) + check(assert_type(read_pickle(path), Any), DataFrame) + + +def test_pickle_compression(): + with ensure_clean() as path: + DF.to_pickle(path, compression="gzip") + check( + assert_type(read_pickle(path, compression="gzip"), Any), + DataFrame, + ) + + check( + assert_type(read_pickle(path, compression="gzip"), Any), + DataFrame, + ) + + +def test_pickle_storage_options(): + with ensure_clean() as path: + DF.to_pickle(path, storage_options={}) + + check( + assert_type(read_pickle(path, storage_options={}), Any), + DataFrame, + ) + + +def test_to_pickle_series(): + s: Series = DF["a"] + with ensure_clean() as path: + check(assert_type(s.to_pickle(path), None), type(None)) + check(assert_type(read_pickle(path), Any), Series) + + def test_read_stata_df(): with ensure_clean() as path: DF.to_stata(path)