ENH/BUG: Add arcname to to_csv for ZIP compressed csv filename (pandas-dev#26023)

Author: drew-heenan

pandas/core/generic.py

@@ -2921,7 +2921,7 @@ def to_csv(self, path_or_buf=None, sep=",", na_rep='', float_format=None,
                mode='w', encoding=None, compression='infer', quoting=None,
                quotechar='"', line_terminator=None, chunksize=None,
                tupleize_cols=None, date_format=None, doublequote=True,
-               escapechar=None, decimal='.'):
+               escapechar=None, decimal='.', arcname=None):
         r"""
         Write object to a comma-separated values (csv) file.
 
@@ -3011,6 +3011,9 @@ def to_csv(self, path_or_buf=None, sep=",", na_rep='', float_format=None,
         decimal : str, default '.'
             Character recognized as decimal separator. E.g. use ',' for
             European data.
+        arcname : str, default None
+            Name of file within a ZIP archive. Only used when `path_or_buf` is
+            a path and `compression` is set to or inferred as 'zip'.
 
         Returns
         -------
@@ -3053,7 +3056,8 @@ def to_csv(self, path_or_buf=None, sep=",", na_rep='', float_format=None,
                                  tupleize_cols=tupleize_cols,
                                  date_format=date_format,
                                  doublequote=doublequote,
-                                 escapechar=escapechar, decimal=decimal)
+                                 escapechar=escapechar, decimal=decimal,
+                                 arcname=arcname)
         formatter.save()
 
         if path_or_buf is None:

pandas/io/common.py

@@ -288,7 +288,7 @@ def _infer_compression(filepath_or_buffer, compression):
 
 
 def _get_handle(path_or_buf, mode, encoding=None, compression=None,
-                memory_map=False, is_text=True):
+                memory_map=False, is_text=True, arcname=None):
     """
     Get file handle for given path/buffer and mode.
 
@@ -350,7 +350,7 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
 
         # ZIP Compression
         elif compression == 'zip':
-            zf = BytesZipFile(path_or_buf, mode)
+            zf = BytesZipFile(path_or_buf, mode, arcname=arcname)
             # Ensure the container is closed as well.
             handles.append(zf)
             if zf.mode == 'w':
@@ -420,13 +420,15 @@ class BytesZipFile(zipfile.ZipFile, BytesIO):  # type: ignore
     bytes strings into a member of the archive.
     """
     # GH 17778
-    def __init__(self, file, mode, compression=zipfile.ZIP_DEFLATED, **kwargs):
+    def __init__(self, file, mode, compression=zipfile.ZIP_DEFLATED, arcname=None, **kwargs):
         if mode in ['wb', 'rb']:
             mode = mode.replace('b', '')
+        self.arcname = arcname
         super(BytesZipFile, self).__init__(file, mode, compression, **kwargs)
 
     def write(self, data):
-        super(BytesZipFile, self).writestr(self.filename, data)
+        arcname = self.filename if self.arcname is None else self.arcname
+        super(BytesZipFile, self).writestr(arcname, data)
 
     @property
     def closed(self):

pandas/io/formats/csvs.py

@@ -29,7 +29,7 @@ def __init__(self, obj, path_or_buf=None, sep=",", na_rep='',
                  compression='infer', quoting=None, line_terminator='\n',
                  chunksize=None, tupleize_cols=False, quotechar='"',
                  date_format=None, doublequote=True, escapechar=None,
-                 decimal='.'):
+                 decimal='.', arcname=None):
 
         self.obj = obj
 
@@ -123,6 +123,8 @@ def __init__(self, obj, path_or_buf=None, sep=",", na_rep='',
         if not index:
             self.nlevels = 0
 
+        self.arcname = arcname
+
     def save(self):
         """
         Create the writer & save
@@ -176,7 +178,8 @@ def save(self):
                 else:
                     f, handles = _get_handle(self.path_or_buf, self.mode,
                                              encoding=self.encoding,
-                                             compression=self.compression)
+                                             compression=self.compression,
+                                             arcname=self.arcname)
                     f.write(buf)
                     close = True
             if close:

pandas/tests/io/formats/test_to_csv.py

@@ -537,3 +537,19 @@ def test_to_csv_compression(self, compression_only,
             result = pd.read_csv(path, index_col=0,
                                  compression=read_compression)
             tm.assert_frame_equal(result, df)
+
+    @pytest.mark.parametrize("compression", ["zip", "infer"])
+    @pytest.mark.parametrize("arcname", [None, "test_to_csv.csv",
+                                         "test_to_csv.zip"])
+    def test_to_csv_zip_arcname(self, compression, arcname):
+        # GH 26023
+        from zipfile import ZipFile
+
+        df = DataFrame({"ABC": [1]})
+        with tm.ensure_clean("to_csv_arcname.zip") as path:
+            df.to_csv(path, compression=compression,
+                      arcname=arcname)
+            zp = ZipFile(path)
+            expected_arcname = path if arcname is None else arcname
+            assert len(zp.filelist) == 1
+            assert zp.filelist[0].filename == expected_arcname