ENH: Allow compression in NDFrame.to_csv to be a dict with optional arguments (#26023)

doc/source/whatsnew/v0.25.0.rst

@@ -202,6 +202,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:`DataFrame.to_csv` and :meth:`Series.to_csv` now support dicts as ``compression`` argument with key ``'method'`` being the compression method and others as additional compression options when the compression method is ``'zip'``. (:issue:`26023`)
 - :func:`merge_asof` now gives a more clear error message when merge keys are categoricals that are not equal (:issue:`26136`)
 - :meth:`pandas.core.window.Rolling` supports exponential (or Poisson) window type (:issue:`21303`)
 - Error message for missing required imports now includes the original import error's text (:issue:`23868`)

pandas/core/generic.py

@@ -6,7 +6,17 @@
 import operator
 import pickle
 from textwrap import dedent
-from typing import Callable, FrozenSet, List, Optional, Set
+from typing import (
+    Callable,
+    Dict,
+    FrozenSet,
+    Hashable,
+    List,
+    Optional,
+    Set,
+    Sequence,
+    Union,
+)
 import warnings
 import weakref
 
@@ -15,6 +25,7 @@
 from pandas._config import config
 
 from pandas._libs import Timestamp, iNaT, properties
+from pandas._typing import FilePathOrBuffer
 from pandas.compat import set_function_name
 from pandas.compat._optional import import_optional_dependency
 from pandas.compat.numpy import function as nv
@@ -121,6 +132,9 @@ def _single_replace(self, to_replace, method, inplace, limit):
     return result
 
 
+bool_t = bool  # Need alias because NDFrame has def bool:
+
+
 class NDFrame(PandasObject, SelectionMixin):
     """
     N-dimensional analogue of DataFrame. Store multi-dimensional in a
@@ -3078,26 +3092,26 @@ def to_latex(
 
     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,
-        quotechar='"',
-        line_terminator=None,
-        chunksize=None,
-        date_format=None,
-        doublequote=True,
-        escapechar=None,
-        decimal=".",
-    ):
+        path_or_buf: Optional[FilePathOrBuffer] = None,
+        sep: str = ",",
+        na_rep: str = "",
+        float_format: Optional[str] = None,
+        columns: Optional[Sequence[Hashable]] = None,
+        header: Union[bool_t, List[str]] = True,
+        index: bool_t = True,
+        index_label: Optional[Union[bool_t, str, Sequence[Hashable]]] = None,
+        mode: str = "w",
+        encoding: Optional[str] = None,
+        compression: Optional[Union[str, Dict[str, str]]] = "infer",
+        quoting: Optional[int] = None,
+        quotechar: str = '"',
+        line_terminator: Optional[str] = None,
+        chunksize: Optional[int] = None,
+        date_format: Optional[str] = None,
+        doublequote: bool_t = True,
+        escapechar: Optional[str] = None,
+        decimal: Optional[str] = ".",
+    ) -> Optional[str]:
         r"""
         Write object to a comma-separated values (csv) file.
 
@@ -3144,16 +3158,21 @@ def to_csv(
         encoding : str, optional
             A string representing the encoding to use in the output file,
             defaults to 'utf-8'.
-        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', other entries passed as
+            additional compression options.
+
+            .. versionchanged:: 0.25.0
+
+               May now be a dict with key 'method' as compression mode
+               and other entries as additional compression options if
+               compression mode is 'zip'.
 
         quoting : optional constant from csv module
             Defaults to csv.QUOTE_MINIMAL. If you have set a `float_format`
@@ -3198,6 +3217,13 @@ def to_csv(
         ...                    'weapon': ['sai', 'bo staff']})
         >>> df.to_csv(index=False)
         'name,mask,weapon\nRaphael,red,sai\nDonatello,purple,bo staff\n'
+
+        # create 'out.zip' containing 'out.csv'
+        >>> compression_opts = dict(method='zip',
+        ...                         archive_name='out.csv')  # doctest: +SKIP
+
+        >>> df.to_csv('out.zip', index=False,
+        ...           compression=compression_opts)  # doctest: +SKIP
         """
 
         df = self if isinstance(self, ABCDataFrame) else self.to_frame()
@@ -3231,6 +3257,8 @@ def to_csv(
         if path_or_buf is None:
             return formatter.path_or_buf.getvalue()
 
+        return None
+
     # ----------------------------------------------------------------------
     # Fancy Indexing
 

pandas/io/common.py

@@ -10,6 +10,7 @@
 import mmap
 import os
 import pathlib
+from typing import Any, Dict, Optional, Tuple, Type, Union
 from urllib.error import URLError  # noqa
 from urllib.parse import (  # noqa
     urlencode,
@@ -22,6 +23,7 @@
 from urllib.request import pathname2url, urlopen
 import zipfile
 
+from pandas._typing import FilePathOrBuffer
 from pandas.errors import (  # noqa
     AbstractMethodError,
     DtypeWarning,
@@ -242,13 +244,46 @@ def file_path_to_url(path):
 _compression_to_extension = {"gzip": ".gz", "bz2": ".bz2", "zip": ".zip", "xz": ".xz"}
 
 
+def _get_compression_method(
+    compression: Optional[Union[str, Dict[str, str]]]
+) -> Tuple[Optional[str], Dict[str, str]]:
+    """
+    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}, Optional[str]
+              {compression arguments}, Dict[str, str])
+
+    Raises
+    ------
+    ValueError on dict missing 'method' key
+    """
+    # Handle dict
+    if isinstance(compression, dict):
+        compression_args = compression.copy()
+        try:
+            compression = compression_args.pop("method")
+        except KeyError:
+            raise ValueError("If dict, compression must have key 'method'")
+    else:
+        compression_args = {}
+    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 :
@@ -257,12 +292,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
@@ -297,7 +330,12 @@ def _infer_compression(filepath_or_buffer, compression):
 
 
 def _get_handle(
-    path_or_buf, mode, encoding=None, compression=None, memory_map=False, is_text=True
+    path_or_buf,
+    mode: str,
+    encoding=None,
+    compression: Optional[Union[str, Dict[str, Any]]] = None,
+    memory_map: bool = False,
+    is_text: bool = True,
 ):
     """
     Get file handle for given path/buffer and mode.
@@ -309,10 +347,21 @@ def _get_handle(
     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', other entries passed as additional compression options.
+
+        .. versionchanged:: 0.25.0
+
+           May now be a dict with key 'method' as compression mode
+           and other keys as compression options if compression
+           mode is 'zip'.
+
     memory_map : boolean, default False
         See parsers._parser_params for more information.
     is_text : boolean, default True
@@ -326,12 +375,13 @@ def _get_handle(
     handles : list of file-like objects
         A list of file-like object that were opened in this function.
     """
+    need_text_wrapping = (BytesIO,)  # type: Tuple[Type[BytesIO], ...]
     try:
         from s3fs import S3File
 
-        need_text_wrapping = (BytesIO, S3File)
+        need_text_wrapping = need_text_wrapping + (S3File,)
     except ImportError:
-        need_text_wrapping = (BytesIO,)
+        pass
 
     handles = list()
     f = path_or_buf
@@ -340,6 +390,7 @@ def _get_handle(
     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)
 
@@ -361,7 +412,7 @@ def _get_handle(
 
         # 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":
@@ -435,13 +486,23 @@ class BytesZipFile(zipfile.ZipFile, BytesIO):  # type: ignore
     """
 
     # GH 17778
-    def __init__(self, file, mode, compression=zipfile.ZIP_DEFLATED, **kwargs):
+    def __init__(
+        self,
+        file: FilePathOrBuffer,
+        mode: str,
+        archive_name: Optional[str] = None,
+        **kwargs
+    ):
         if mode in ["wb", "rb"]:
             mode = mode.replace("b", "")
-        super().__init__(file, mode, compression, **kwargs)
+        self.archive_name = archive_name
+        super().__init__(file, mode, zipfile.ZIP_DEFLATED, **kwargs)
 
     def write(self, data):
-        super().writestr(self.filename, data)
+        archive_name = self.filename
+        if self.archive_name is not None:
+            archive_name = self.archive_name
+        super().writestr(archive_name, data)
 
     @property
     def closed(self):

pandas/io/formats/csvs.py

@@ -22,6 +22,7 @@
 
 from pandas.io.common import (
     UnicodeWriter,
+    _get_compression_method,
     _get_handle,
     _infer_compression,
     get_filepath_or_buffer,
@@ -58,6 +59,9 @@ def __init__(
         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
         )
@@ -180,7 +184,7 @@ def save(self):
                 self.path_or_buf,
                 self.mode,
                 encoding=self.encoding,
-                compression=self.compression,
+                compression=dict(self.compression_args, method=self.compression),
             )
             close = True
 
@@ -208,11 +212,13 @@ 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