BUG: GH17778 - DataFrame.to_pickle() fails for .zip format.

Author: kchomski

pandas/io/common.py

@@ -357,17 +357,20 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
         # ZIP Compression
         elif compression == 'zip':
             import zipfile
-            zip_file = zipfile.ZipFile(path_or_buf)
-            zip_names = zip_file.namelist()
-            if len(zip_names) == 1:
-                f = zip_file.open(zip_names.pop())
-            elif len(zip_names) == 0:
-                raise ValueError('Zero files found in ZIP file {}'
-                                 .format(path_or_buf))
+            if mode == 'wb':
+                f = zipfile.ZipFile(path_or_buf, 'w')
             else:
-                raise ValueError('Multiple files found in ZIP file.'
-                                 ' Only one file per ZIP: {}'
-                                 .format(zip_names))
+                zip_file = zipfile.ZipFile(path_or_buf)
+                zip_names = zip_file.namelist()
+                if len(zip_names) == 1:
+                    f = zip_file.open(zip_names.pop())
+                elif len(zip_names) == 0:
+                    raise ValueError('Zero files found in ZIP file {}'
+                                     .format(path_or_buf))
+                else:
+                    raise ValueError('Multiple files found in ZIP file.'
+                                     ' Only one file per ZIP: {}'
+                                     .format(zip_names))
 
         # XZ Compression
         elif compression == 'xz':

pandas/io/pickle.py

@@ -42,7 +42,15 @@ def to_pickle(obj, path, compression='infer', protocol=pkl.HIGHEST_PROTOCOL):
     if protocol < 0:
         protocol = pkl.HIGHEST_PROTOCOL
     try:
-        pkl.dump(obj, f, protocol=protocol)
+        import zipfile
+        if isinstance(f, zipfile.ZipFile):
+            import tempfile
+            tmp_file = tempfile.NamedTemporaryFile()
+            pkl.dump(obj, tmp_file, protocol=protocol)
+            tmp_file.seek(0)
+            f.write(tmp_file.name)
+        else:
+            pkl.dump(obj, f, protocol=protocol)
     finally:
         for _f in fh:
             _f.close()