Repository: pandas-dev/pandas
Pull request:  ENH: Allow compression in NDFrame.to_csv to be a dict with optional arguments (#26023)
- [x] closes #26023 
- [x] tests added / passed
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry


Problem statement:
pd.DataFrame.to_csv('filename.zip') doesn't extract with a '.csv' extension
#### Code Sample, a copy-pastable example if possible

```python
df = pd.DataFrame({'name': ['Raphael', 'Donatello'], 'mask': ['red', 'purple']})

# When trying to create a compressed csv, these give odd results.
df.to_csv('out.csv', compression='zip')  # --> zip file named 'out.csv' containing csv file 'out.csv'
df.to_csv('out.zip')                     # --> zip file named 'out.zip' containing csv file 'out.zip'
df.to_csv('out.csv.zip')                 # --> zip file named 'out.csv.zip' containing csv file 'out.csv.zip'

# This would be the desired behaviour, if we had an 'arcname' 
# parameter like zipfile.ZipFile.write(arcname, data)
df.to_csv('out.zip', arcname='data.csv') # --> zip file named 'out.zip' containing csv file 'data.csv'
```
#### Problem description

When `pd.DataFrame.to_csv` creates compressed zip files, the name of the csv file inside the archive is always the same as the name of the zip archive file itself. This is obviously problematic because the archive has a `.zip` extension but we want the csv file to have a `.csv` extension when it is extracted.

Other compression methods meant for a single file like 'bz2', 'gzip', and 'xz' do not have this problem because a file 'file.csv.gz' for instance, will automatically become 'file.csv' when decompressed.

This would be a relatively easy fix by adding an `arcname=None` parameter to `to_csv`, passing it through `pandas.io.formats.csvs.CSVFormatter` to `pandas.io.formats.csvs._get_handle` and using that instead of `ZipFile.filename` if provided.

#### Expected Output

See comments in Code Sample above for expected output.
 
#### Output of ``pd.show_versions()``

<details>

INSTALLED VERSIONS
------------------
commit: None
python: 3.7.2.final.0
python-bits: 64
OS: Linux
OS-release: 4.18.0-17-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: en_US.UTF-8

pandas: 0.24.1
pytest: 4.3.0
pip: 19.0.3
setuptools: 40.0.0
Cython: 0.28.5
numpy: 1.16.2
scipy: 1.2.1
pyarrow: 0.11.1
xarray: None
IPython: 7.1.1
sphinx: 1.8.5
patsy: 0.5.1
dateutil: 2.7.3
pytz: 2018.5
blosc: None
bottleneck: None
tables: None
numexpr: None
feather: None
matplotlib: 3.0.3
openpyxl: 2.4.11
xlrd: 1.1.0
xlwt: None
xlsxwriter: None
lxml.etree: None
bs4: 4.6.3
html5lib: 1.0.1
sqlalchemy: 1.2.11
pymysql: None
psycopg2: None
jinja2: 2.10
s3fs: None
fastparquet: 0.2.1
pandas_gbq: None
pandas_datareader: None
gcsfs: None
</details>

Hints from issue discussion:
I'll take a look at this!

Commit message: compression args passed as kwargs, update relevant docs

Patch to review:
diff --git a/doc/source/whatsnew/v0.25.0.rst b/doc/source/whatsnew/v0.25.0.rst
index fe047b4a141ef..2b50ffeac7859 100644
--- a/doc/source/whatsnew/v0.25.0.rst
+++ b/doc/source/whatsnew/v0.25.0.rst
@@ -36,6 +36,7 @@ Other Enhancements
 - :class:`RangeIndex` has gained :attr:`~RangeIndex.start`, :attr:`~RangeIndex.stop`, and :attr:`~RangeIndex.step` attributes (:issue:`25710`)
 - :class:`datetime.timezone` objects are now supported as arguments to timezone methods and constructors (:issue:`25065`)
 - :meth:`DataFrame.query` and :meth:`DataFrame.eval` now supports quoting column names with backticks to refer to names with spaces (:issue:`6508`)
+- :meth:`NDFrame.to_csv` now supports dicts as ``compression`` argument with key ``'method'`` being the compression method and others as kwargs of ``ByteZipFile`` when the compression method is ``'zip'``. (:issue:`26023`)
 
 .. _whatsnew_0250.api_breaking:
 
diff --git a/pandas/core/generic.py b/pandas/core/generic.py
index 5e8876404d1ad..6bd6809689f49 100644
--- a/pandas/core/generic.py
+++ b/pandas/core/generic.py
@@ -7,6 +7,7 @@
 import operator
 import pickle
 from textwrap import dedent
+from typing import Dict, Union
 import warnings
 import weakref
 
@@ -2918,7 +2919,8 @@ def to_latex(self, buf=None, columns=None, col_space=None, header=True,
 
     def to_csv(self, path_or_buf=None, sep=",", na_rep='', float_format=None,
                columns=None, header=True, index=True, index_label=None,
-               mode='w', encoding=None, compression='infer', quoting=None,
+               mode='w', encoding=None,
+               compression: Union[str, Dict, None] = 'infer', quoting=None,
                quotechar='"', line_terminator=None, chunksize=None,
                tupleize_cols=None, date_format=None, doublequote=True,
                escapechar=None, decimal='.'):
@@ -2968,16 +2970,22 @@ def to_csv(self, path_or_buf=None, sep=",", na_rep='', float_format=None,
         encoding : str, optional
             A string representing the encoding to use in the output file,
             defaults to 'ascii' on Python 2 and 'utf-8' on Python 3.
-        compression : str, default 'infer'
-            Compression mode among the following possible values: {'infer',
-            'gzip', 'bz2', 'zip', 'xz', None}. If 'infer' and `path_or_buf`
-            is path-like, then detect compression from the following
-            extensions: '.gz', '.bz2', '.zip' or '.xz'. (otherwise no
-            compression).
-
-            .. versionchanged:: 0.24.0
-
-               'infer' option added and set to default.
+        compression : str or dict, default 'infer'
+            If str, represents compression mode. If dict, value at 'method' is
+            the compression mode. Compression mode may be any of the following
+            possible values: {'infer', 'gzip', 'bz2', 'zip', 'xz', None}. If
+            compression mode is 'infer' and `path_or_buf` is path-like, then
+            detect compression mode from the following extensions: '.gz',
+            '.bz2', '.zip' or '.xz'. (otherwise no compression). If dict given
+            and mode is 'zip' or inferred as 'zip', optional value at 'arcname'
+            specifies name of file within ZIP archive, assuming equal to
+            `path_or_buf` if not specified or None.
+
+            .. versionchanged:: 0.25.0
+
+               May now be a dict with key 'method' as compression mode
+               and other entries as ByteZipFile kwargs if compression mode
+               is 'zip'
 
         quoting : optional constant from csv module
             Defaults to csv.QUOTE_MINIMAL. If you have set a `float_format`
diff --git a/pandas/io/common.py b/pandas/io/common.py
index bcbfa7930311e..8ee8122f017b6 100644
--- a/pandas/io/common.py
+++ b/pandas/io/common.py
@@ -9,6 +9,7 @@
 import lzma
 import mmap
 import os
+from typing import Dict, Tuple, Union
 from urllib.error import URLError  # noqa
 from urllib.parse import (  # noqa
     urlencode, urljoin, urlparse as parse_url, uses_netloc, uses_params,
@@ -233,13 +234,45 @@ def file_path_to_url(path):
 }
 
 
+def _get_compression_method(compression: Union[str, Dict, None]):
+    """
+    Simplifies a compression argument to a compression method string and
+    a dict containing additional arguments.
+
+    Parameters
+    ----------
+    compression : str or dict
+        If string, specifies the compression method. If dict, value at key
+        'method' specifies compression method.
+
+    Returns
+    -------
+    tuple of ({compression method}, any
+              {compression arguments}, dict)
+
+    Raises
+    ------
+    ValueError on dict missing 'method' key
+    """
+    compression_args = {}  # type: Dict
+    # Handle dict
+    if isinstance(compression, dict):
+        compression_args = compression.copy()
+        try:
+            compression = compression['method']
+            compression_args.pop('method')
+        except KeyError:
+            raise ValueError("If dict, compression "
+                             "must have key 'method'")
+    return compression, compression_args
+
+
 def _infer_compression(filepath_or_buffer, compression):
     """
     Get the compression method for filepath_or_buffer. If compression='infer',
     the inferred compression method is returned. Otherwise, the input
     compression method is returned unchanged, unless it's invalid, in which
     case an error is raised.
-
     Parameters
     ----------
     filepath_or_buffer :
@@ -248,12 +281,10 @@ def _infer_compression(filepath_or_buffer, compression):
         If 'infer' and `filepath_or_buffer` is path-like, then detect
         compression from the following extensions: '.gz', '.bz2', '.zip',
         or '.xz' (otherwise no compression).
-
     Returns
     -------
     string or None :
         compression method
-
     Raises
     ------
     ValueError on invalid compression specified
@@ -287,8 +318,9 @@ def _infer_compression(filepath_or_buffer, compression):
     raise ValueError(msg)
 
 
-def _get_handle(path_or_buf, mode, encoding=None, compression=None,
-                memory_map=False, is_text=True):
+def _get_handle(path_or_buf, mode, encoding=None,
+                compression: Union[str, Dict, None] = None, memory_map=False,
+                is_text=True):
     """
     Get file handle for given path/buffer and mode.
 
@@ -299,10 +331,22 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
     mode : str
         mode to open path_or_buf with
     encoding : str or None
-    compression : {'infer', 'gzip', 'bz2', 'zip', 'xz', None}, default None
-        If 'infer' and `filepath_or_buffer` is path-like, then detect
-        compression from the following extensions: '.gz', '.bz2', '.zip',
-        or '.xz' (otherwise no compression).
+    compression : str or dict, default None
+        If string, specifies compression mode. If dict, value at key 'method'
+        specifies compression mode. Compression mode must be one of {'infer',
+        'gzip', 'bz2', 'zip', 'xz', None}. If compression mode is 'infer'
+        and `filepath_or_buffer` is path-like, then detect compression from
+        the following extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise
+        no compression). If dict and compression mode is 'zip' or inferred as
+        'zip', optional value at key 'arcname' specifies the name of the file
+        within ZIP archive at `path_or_buf`.
+
+        .. versionchanged:: 0.25.0
+
+           May now be a dict with key 'method' as compression mode
+           and other keys as kwargs for ByteZipFile if compression
+           mode is 'zip'.
+
     memory_map : boolean, default False
         See parsers._parser_params for more information.
     is_text : boolean, default True
@@ -318,7 +362,7 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
     """
     try:
         from s3fs import S3File
-        need_text_wrapping = (BytesIO, S3File)
+        need_text_wrapping = (BytesIO, S3File)  # type: Tuple
     except ImportError:
         need_text_wrapping = (BytesIO,)
 
@@ -329,6 +373,7 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
     path_or_buf = _stringify_path(path_or_buf)
     is_path = isinstance(path_or_buf, str)
 
+    compression, compression_args = _get_compression_method(compression)
     if is_path:
         compression = _infer_compression(path_or_buf, compression)
 
@@ -350,7 +395,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, **compression_args)
             # Ensure the container is closed as well.
             handles.append(zf)
             if zf.mode == 'w':
@@ -420,13 +465,16 @@ 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: Union[str, zipfile.ZipInfo, None] = None, **kwargs):
         if mode in ['wb', 'rb']:
             mode = mode.replace('b', '')
-        super(BytesZipFile, self).__init__(file, mode, compression, **kwargs)
+        self.arcname = arcname
+        super().__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().writestr(arcname, data)
 
     @property
     def closed(self):
diff --git a/pandas/io/formats/csvs.py b/pandas/io/formats/csvs.py
index 37f1372366545..2fceac352faed 100644
--- a/pandas/io/formats/csvs.py
+++ b/pandas/io/formats/csvs.py
@@ -6,6 +6,7 @@
 import csv as csvlib
 from io import StringIO
 import os
+from typing import Dict, Union
 import warnings
 from zipfile import ZipFile
 
@@ -18,7 +19,8 @@
 from pandas.core.dtypes.missing import notna
 
 from pandas.io.common import (
-    UnicodeWriter, _get_handle, _infer_compression, get_filepath_or_buffer)
+    UnicodeWriter, _get_compression_method, _get_handle, _infer_compression,
+    get_filepath_or_buffer)
 
 
 class CSVFormatter(object):
@@ -26,18 +28,23 @@ class CSVFormatter(object):
     def __init__(self, obj, path_or_buf=None, sep=",", na_rep='',
                  float_format=None, cols=None, header=True, index=True,
                  index_label=None, mode='w', nanRep=None, encoding=None,
-                 compression='infer', quoting=None, line_terminator='\n',
-                 chunksize=None, tupleize_cols=False, quotechar='"',
-                 date_format=None, doublequote=True, escapechar=None,
-                 decimal='.'):
+                 compression: Union[str, Dict, None] = 'infer', quoting=None,
+                 line_terminator='\n', chunksize=None, tupleize_cols=False,
+                 quotechar='"', date_format=None, doublequote=True,
+                 escapechar=None, decimal='.'):
 
         self.obj = obj
 
         if path_or_buf is None:
             path_or_buf = StringIO()
 
+        # Extract compression mode as given, if dict
+        compression, self.compression_args \
+            = _get_compression_method(compression)
+
         self.path_or_buf, _, _, _ = get_filepath_or_buffer(
-            path_or_buf, encoding=encoding, compression=compression, mode=mode
+            path_or_buf, encoding=encoding,
+            compression=compression, mode=mode
         )
         self.sep = sep
         self.na_rep = na_rep
@@ -115,7 +122,7 @@ def __init__(self, obj, path_or_buf=None, sep=",", na_rep='',
         self.data_index = obj.index
         if (isinstance(self.data_index, (ABCDatetimeIndex, ABCPeriodIndex)) and
                 date_format is not None):
-            from pandas import Index
+            from pandas import Index  # type: ignore
             self.data_index = Index([x.strftime(date_format) if notna(x) else
                                      '' for x in self.data_index])
 
@@ -150,7 +157,8 @@ def save(self):
         else:
             f, handles = _get_handle(self.path_or_buf, self.mode,
                                      encoding=self.encoding,
-                                     compression=self.compression)
+                                     compression=dict(self.compression_args,
+                                                      method=self.compression))
             close = True
 
         try:
@@ -174,9 +182,11 @@ def save(self):
                 if hasattr(self.path_or_buf, 'write'):
                     self.path_or_buf.write(buf)
                 else:
+                    compression = dict(self.compression_args,
+                                       method=self.compression)
                     f, handles = _get_handle(self.path_or_buf, self.mode,
                                              encoding=self.encoding,
-                                             compression=self.compression)
+                                             compression=compression)
                     f.write(buf)
                     close = True
             if close:
diff --git a/pandas/tests/io/formats/test_to_csv.py b/pandas/tests/io/formats/test_to_csv.py
index fbd71dfa8262b..1592363f34190 100644
--- a/pandas/tests/io/formats/test_to_csv.py
+++ b/pandas/tests/io/formats/test_to_csv.py
@@ -537,3 +537,32 @@ 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)
+
+    def test_to_csv_compression_dict(self, compression_only):
+        # GH 26023
+        method = compression_only
+        df = DataFrame({"ABC": [1]})
+        filename = "to_csv_compress_as_dict."
+        filename += "gz" if method == "gzip" else method
+        with tm.ensure_clean(filename) as path:
+            df.to_csv(path, compression={"method": method})
+            read_df = pd.read_csv(path, index_col=0)
+            tm.assert_frame_equal(read_df, 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={"method": compression,
+                                         "arcname": arcname})
+            zp = ZipFile(path)
+            expected_arcname = path if arcname is None else arcname
+            expected_arcname = os.path.basename(expected_arcname)
+            assert len(zp.filelist) == 1
+            archived_file = os.path.basename(zp.filelist[0].filename)
+            assert archived_file == expected_arcname


Code review comments:
diff_hunk: @@ -318,7 +362,7 @@ def _get_handle(path_or_buf, mode, encoding=None, compression=None,
     """
     try:
         from s3fs import S3File
-        need_text_wrapping = (BytesIO, S3File)
+        need_text_wrapping = (BytesIO, S3File)  # type: Tuple
path: pandas/io/common.py
text: @user1:
Was this throwing a Typing error? Think MyPy should be able to infer here

@author:
@user1 Yeah, it was. Couldn’t tell why, but MyPy couldn’t infer when I added types to the function definition. It also occurred on `from pandas import Index` in `csvs.py`

@user1:
What error was it giving you?

@author:
It was giving `pandas/io/common.py:367: error: Incompatible types in assignment (expression has type "Tuple[Type[BytesIO]]", variable has type "Tuple[Type[BytesIO], Any]")` only if any of the function's arguments have type annotations.

diff_hunk: @@ -233,13 +234,45 @@ def file_path_to_url(path):
 }
 
 
+def _get_compression_method(compression: Union[str, Dict, None]):
+    """
+    Simplifies a compression argument to a compression method string and
+    a dict containing additional arguments.
+
+    Parameters
+    ----------
+    compression : str or dict
+        If string, specifies the compression method. If dict, value at key
+        'method' specifies compression method.
+
+    Returns
+    -------
+    tuple of ({compression method}, any
+              {compression arguments}, dict)
+
+    Raises
+    ------
+    ValueError on dict missing 'method' key
+    """
+    compression_args = {}  # type: Dict
path: pandas/io/common.py
text: @user1:
Can you add the key / value types for the Dict

diff_hunk: @@ -18,26 +19,32 @@
 from pandas.core.dtypes.missing import notna
 
 from pandas.io.common import (
-    UnicodeWriter, _get_handle, _infer_compression, get_filepath_or_buffer)
+    UnicodeWriter, _get_compression_method, _get_handle, _infer_compression,
+    get_filepath_or_buffer)
 
 
 class CSVFormatter(object):
 
     def __init__(self, obj, path_or_buf=None, sep=",", na_rep='',
                  float_format=None, cols=None, header=True, index=True,
                  index_label=None, mode='w', nanRep=None, encoding=None,
-                 compression='infer', quoting=None, line_terminator='\n',
-                 chunksize=None, tupleize_cols=False, quotechar='"',
-                 date_format=None, doublequote=True, escapechar=None,
-                 decimal='.'):
+                 compression: Union[str, Dict, None] = 'infer', quoting=None,
+                 line_terminator='\n', chunksize=None, tupleize_cols=False,
+                 quotechar='"', date_format=None, doublequote=True,
+                 escapechar=None, decimal='.'):
 
         self.obj = obj
 
         if path_or_buf is None:
             path_or_buf = StringIO()
 
+        # Extract compression mode as given, if dict
+        compression, self.compression_args \
path: pandas/io/formats/csvs.py
text: @user1:
Can you break at the function signature instead? We don't use explicit line continuations like this in code base

diff_hunk: @@ -287,8 +318,9 @@ def _infer_compression(filepath_or_buffer, compression):
     raise ValueError(msg)
 
 
-def _get_handle(path_or_buf, mode, encoding=None, compression=None,
-                memory_map=False, is_text=True):
+def _get_handle(path_or_buf, mode, encoding=None,
+                compression: Union[str, Dict, None] = None, memory_map=False,
path: pandas/io/common.py
text: @user1:
Same here for breaking up signature