ENH: Synchronize pickle with upstream (#206)

* ENH: Synchronize pickle with upstream

Align APIs
Add tests

* BUG: Correct ReadPickleBuffer

Correct definition
Use ensure_clean

* MAINT: Clean up to_pickle testing

Ensure to_pickle is tested since it  appears in pandas/io/api

* BUG: Restore read_pickle

* CLN: Remove to_pickle from frame

Remove extra def
Test on series

* MAINT: Remove extra def

Co-authored-by: Kevin Sheppard <kevin.sheppard@gmail.com>

Author: bashtage

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: ...

pandas-stubs/core/frame.pyi

@@ -1935,12 +1935,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,

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

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: ...

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)