Skip to content

ENH: Synchronize pickle with upstream #206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Sep 5, 2022
Merged
3 changes: 3 additions & 0 deletions pandas-stubs/_typing.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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: ...
Expand Down
6 changes: 0 additions & 6 deletions pandas-stubs/core/frame.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion pandas-stubs/core/generic.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ from pandas._typing import (
SortKind,
StorageOptions,
T,
WriteBuffer,
)

from pandas.io.pytables import HDFStore
Expand Down Expand Up @@ -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
Expand Down
21 changes: 14 additions & 7 deletions pandas-stubs/io/pickle.pyi
Original file line number Diff line number Diff line change
@@ -1,15 +1,22 @@
from typing import Any

from pandas._typing import (
CompressionOptions,
FilePathOrBuffer,
FilePath,
ReadPickleBuffer,
StorageOptions,
WriteBuffer,
)

def to_pickle(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since pandas.io.pickle.to_pickle() is not public, we should delete this here.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

still need to handle removal of to_pickle() from this file (and then the associated tests)

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: ...
64 changes: 64 additions & 0 deletions tests/test_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import sqlite3
from typing import (
TYPE_CHECKING,
Any,
Dict,
Generator,
List,
Expand All @@ -30,6 +31,7 @@
read_json,
read_orc,
read_parquet,
read_pickle,
read_sas,
read_spss,
read_sql,
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down