ENH: Add compression to stata writers (#34013)

doc/source/whatsnew/v1.1.0.rst

@@ -227,7 +227,12 @@ Other enhancements
 - The ``ExtensionArray`` class has now an :meth:`~pandas.arrays.ExtensionArray.equals`
   method, similarly to :meth:`Series.equals` (:issue:`27081`).
 - The minimum suppported dta version has increased to 105 in :meth:`~pandas.io.stata.read_stata` and :class:`~pandas.io.stata.StataReader`  (:issue:`26667`).
-
+- :meth:`~pandas.core.frame.DataFrame.to_stata` supports compression using the ``compression``
+  keyword argument. Compression can either be inferred or explicitly set using a string or a
+  dictionary containing both the method and any additional arguments that are passed to the
+  compression library. Compression was also added to the low-level Stata-file writers
+  :class:`~pandas.io.stata.StataWriter`, :class:`~pandas.io.stata.StataWriter117`,
+  and :class:`~pandas.io.stata.StataWriterUTF8` (:issue:`26599`).
 
 .. ---------------------------------------------------------------------------
 

pandas/core/frame.py

@@ -25,6 +25,7 @@
     Iterable,
     Iterator,
     List,
+    Mapping,
     Optional,
     Sequence,
     Set,
@@ -2015,6 +2016,7 @@ def to_stata(
         variable_labels: Optional[Dict[Label, str]] = None,
         version: Optional[int] = 114,
         convert_strl: Optional[Sequence[Label]] = None,
+        compression: Union[str, Mapping[str, str], None] = "infer",
     ) -> None:
         """
         Export DataFrame object to Stata dta format.
@@ -2078,6 +2080,19 @@ def to_stata(
 
             .. versionadded:: 0.23.0
 
+        compression : str or dict, default 'infer'
+            For on-the-fly compression of the output dta. 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
+            `fname` is path-like, then detect compression from the following
+            extensions: '.gz', '.bz2', '.zip', or '.xz' (otherwise no
+            compression). If dict and compression mode is one of {'zip',
+            'gzip', 'bz2'}, or inferred as one of the above, other entries
+            passed as additional compression options.
+
+            .. versionadded:: 1.1.0
+
         Raises
         ------
         NotImplementedError
@@ -2133,6 +2148,7 @@ def to_stata(
             data_label=data_label,
             write_index=write_index,
             variable_labels=variable_labels,
+            compression=compression,
             **kwargs,
         )
         writer.write_file()

pandas/io/stata.py

@@ -16,7 +16,18 @@
 from pathlib import Path
 import struct
 import sys
-from typing import Any, AnyStr, BinaryIO, Dict, List, Optional, Sequence, Tuple, Union
+from typing import (
+    Any,
+    AnyStr,
+    BinaryIO,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 import warnings
 
 from dateutil.relativedelta import relativedelta
@@ -47,7 +58,13 @@
 from pandas.core.indexes.base import Index
 from pandas.core.series import Series
 
-from pandas.io.common import get_filepath_or_buffer, stringify_path
+from pandas.io.common import (
+    get_compression_method,
+    get_filepath_or_buffer,
+    get_handle,
+    infer_compression,
+    stringify_path,
+)
 
 _version_error = (
     "Version of given Stata file is {version}. pandas supports importing "
@@ -1854,13 +1871,18 @@ def read_stata(
     return data
 
 
-def _open_file_binary_write(fname: FilePathOrBuffer) -> Tuple[BinaryIO, bool]:
+def _open_file_binary_write(
+    fname: FilePathOrBuffer, compression: Union[str, Mapping[str, str], None],
+) -> Tuple[BinaryIO, bool, Optional[Union[str, Mapping[str, str]]]]:
     """
     Open a binary file or no-op if file-like.
 
     Parameters
     ----------
     fname : string path, path object or buffer
+        The file name or buffer.
+    compression : {str, dict, None}
+        The compression method to use.
 
     Returns
     -------
@@ -1871,9 +1893,21 @@ def _open_file_binary_write(fname: FilePathOrBuffer) -> Tuple[BinaryIO, bool]:
     """
     if hasattr(fname, "write"):
         # See https://github.com/python/mypy/issues/1424 for hasattr challenges
-        return fname, False  # type: ignore
+        return fname, False, None  # type: ignore
     elif isinstance(fname, (str, Path)):
-        return open(fname, "wb"), True
+        # Extract compression mode as given, if dict
+        compression_typ, compression_args = get_compression_method(compression)
+        compression_typ = infer_compression(fname, compression_typ)
+        path_or_buf, _, compression_typ, _ = get_filepath_or_buffer(
+            fname, compression=compression_typ
+        )
+        if compression_typ is not None:
+            compression = compression_args
+            compression["method"] = compression_typ
+        else:
+            compression = None
+        f, _ = get_handle(path_or_buf, "wb", compression=compression, is_text=False)
+        return f, True, compression
     else:
         raise TypeError("fname must be a binary file, buffer or path-like.")
 
@@ -2050,6 +2084,17 @@ class StataWriter(StataParser):
     variable_labels : dict
         Dictionary containing columns as keys and variable labels as values.
         Each label must be 80 characters or smaller.
+    compression : str or dict, default 'infer'
+        For on-the-fly compression of the output dta. 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 `fname` is path-like,
+        then detect compression from the following extensions: '.gz', '.bz2',
+        '.zip', or '.xz' (otherwise no compression). If dict and compression
+        mode is one of {'zip', 'gzip', 'bz2'}, or inferred as one of the above,
+        other entries passed as additional compression options.
+
+        .. versionadded:: 1.1.0
 
     Returns
     -------
@@ -2074,7 +2119,12 @@ class StataWriter(StataParser):
     >>> writer = StataWriter('./data_file.dta', data)
     >>> writer.write_file()
 
-    Or with dates
+    Directly write a zip file
+    >>> compression = {"method": "zip", "archive_name": "data_file.dta"}
+    >>> writer = StataWriter('./data_file.zip', data, compression=compression)
+    >>> writer.write_file()
+
+    Save a DataFrame with dates
     >>> from datetime import datetime
     >>> data = pd.DataFrame([[datetime(2000,1,1)]], columns=['date'])
     >>> writer = StataWriter('./date_data_file.dta', data, {'date' : 'tw'})
@@ -2094,6 +2144,7 @@ def __init__(
         time_stamp: Optional[datetime.datetime] = None,
         data_label: Optional[str] = None,
         variable_labels: Optional[Dict[Label, str]] = None,
+        compression: Union[str, Mapping[str, str], None] = "infer",
     ):
         super().__init__()
         self._convert_dates = {} if convert_dates is None else convert_dates
@@ -2102,6 +2153,8 @@ def __init__(
         self._data_label = data_label
         self._variable_labels = variable_labels
         self._own_file = True
+        self._compression = compression
+        self._output_file: Optional[BinaryIO] = None
         # attach nobs, nvars, data, varlist, typlist
         self._prepare_pandas(data)
 
@@ -2389,7 +2442,12 @@ def _encode_strings(self) -> None:
                     self.data[col] = encoded
 
     def write_file(self) -> None:
-        self._file, self._own_file = _open_file_binary_write(self._fname)
+        self._file, self._own_file, compression = _open_file_binary_write(
+            self._fname, self._compression
+        )
+        if compression is not None:
+            self._output_file = self._file
+            self._file = BytesIO()
         try:
             self._write_header(data_label=self._data_label, time_stamp=self._time_stamp)
             self._write_map()
@@ -2434,6 +2492,12 @@ def _close(self) -> None:
         """
         # Some file-like objects might not support flush
         assert self._file is not None
+        if self._output_file is not None:
+            assert isinstance(self._file, BytesIO)
+            bio = self._file
+            bio.seek(0)
+            self._file = self._output_file
+            self._file.write(bio.read())
         try:
             self._file.flush()
         except AttributeError:
@@ -2898,6 +2962,17 @@ class StataWriter117(StataWriter):
         Smaller columns can be converted by including the column name.  Using
         StrLs can reduce output file size when strings are longer than 8
         characters, and either frequently repeated or sparse.
+    compression : str or dict, default 'infer'
+        For on-the-fly compression of the output dta. 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 `fname` is path-like,
+        then detect compression from the following extensions: '.gz', '.bz2',
+        '.zip', or '.xz' (otherwise no compression). If dict and compression
+        mode is one of {'zip', 'gzip', 'bz2'}, or inferred as one of the above,
+        other entries passed as additional compression options.
+
+        .. versionadded:: 1.1.0
 
     Returns
     -------
@@ -2923,8 +2998,12 @@ class StataWriter117(StataWriter):
     >>> writer = StataWriter117('./data_file.dta', data)
     >>> writer.write_file()
 
-    Or with long strings stored in strl format
+    Directly write a zip file
+    >>> compression = {"method": "zip", "archive_name": "data_file.dta"}
+    >>> writer = StataWriter117('./data_file.zip', data, compression=compression)
+    >>> writer.write_file()
 
+    Or with long strings stored in strl format
     >>> data = pd.DataFrame([['A relatively long string'], [''], ['']],
     ...                     columns=['strls'])
     >>> writer = StataWriter117('./data_file_with_long_strings.dta', data,
@@ -2946,6 +3025,7 @@ def __init__(
         data_label: Optional[str] = None,
         variable_labels: Optional[Dict[Label, str]] = None,
         convert_strl: Optional[Sequence[Label]] = None,
+        compression: Union[str, Mapping[str, str], None] = "infer",
     ):
         # Copy to new list since convert_strl might be modified later
         self._convert_strl: List[Label] = []
@@ -2961,6 +3041,7 @@ def __init__(
             time_stamp=time_stamp,
             data_label=data_label,
             variable_labels=variable_labels,
+            compression=compression,
         )
         self._map: Dict[str, int] = {}
         self._strl_blob = b""
@@ -3281,6 +3362,17 @@ class StataWriterUTF8(StataWriter117):
         The dta version to use. By default, uses the size of data to determine
         the version. 118 is used if data.shape[1] <= 32767, and 119 is used
         for storing larger DataFrames.
+    compression : str or dict, default 'infer'
+        For on-the-fly compression of the output dta. 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 `fname` is path-like,
+        then detect compression from the following extensions: '.gz', '.bz2',
+        '.zip', or '.xz' (otherwise no compression). If dict and compression
+        mode is one of {'zip', 'gzip', 'bz2'}, or inferred as one of the above,
+        other entries passed as additional compression options.
+
+        .. versionadded:: 1.1.0
 
     Returns
     -------
@@ -3308,6 +3400,11 @@ class StataWriterUTF8(StataWriter117):
     >>> writer = StataWriterUTF8('./data_file.dta', data)
     >>> writer.write_file()
 
+    Directly write a zip file
+    >>> compression = {"method": "zip", "archive_name": "data_file.dta"}
+    >>> writer = StataWriterUTF8('./data_file.zip', data, compression=compression)
+    >>> writer.write_file()
+
     Or with long strings stored in strl format
 
     >>> data = pd.DataFrame([['ᴀ relatively long ŝtring'], [''], ['']],
@@ -3331,6 +3428,7 @@ def __init__(
         variable_labels: Optional[Dict[Label, str]] = None,
         convert_strl: Optional[Sequence[Label]] = None,
         version: Optional[int] = None,
+        compression: Union[str, Mapping[str, str], None] = "infer",
     ):
         if version is None:
             version = 118 if data.shape[1] <= 32767 else 119
@@ -3352,6 +3450,7 @@ def __init__(
             data_label=data_label,
             variable_labels=variable_labels,
             convert_strl=convert_strl,
+            compression=compression,
         )
         # Override version set in StataWriter117 init
         self._dta_version = version

pandas/tests/io/test_stata.py

@@ -1,10 +1,13 @@
+import bz2
 import datetime as dt
 from datetime import datetime
 import gzip
 import io
+import lzma
 import os
 import struct
 import warnings
+import zipfile
 
 import numpy as np
 import pytest
@@ -1863,3 +1866,60 @@ def test_backward_compat(version, datapath):
     expected = pd.read_stata(ref)
     old_dta = pd.read_stata(old)
     tm.assert_frame_equal(old_dta, expected, check_dtype=False)
+
+
+@pytest.mark.parametrize("version", [114, 117, 118, 119, None])
+@pytest.mark.parametrize("use_dict", [True, False])
+@pytest.mark.parametrize("infer", [True, False])
+def test_compression(compression, version, use_dict, infer):
+    file_name = "dta_inferred_compression.dta"
+    if compression:
+        file_ext = "gz" if compression == "gzip" and not use_dict else compression
+        file_name += f".{file_ext}"
+    compression_arg = compression
+    if infer:
+        compression_arg = "infer"
+    if use_dict:
+        compression_arg = {"method": compression}
+
+    df = DataFrame(np.random.randn(10, 2), columns=list("AB"))
+    df.index.name = "index"
+    with tm.ensure_clean(file_name) as path:
+        df.to_stata(path, version=version, compression=compression_arg)
+        if compression == "gzip":
+            with gzip.open(path, "rb") as comp:
+                fp = io.BytesIO(comp.read())
+        elif compression == "zip":
+            with zipfile.ZipFile(path, "r") as comp:
+                fp = io.BytesIO(comp.read(comp.filelist[0]))
+        elif compression == "bz2":
+            with bz2.open(path, "rb") as comp:
+                fp = io.BytesIO(comp.read())
+        elif compression == "xz":
+            with lzma.open(path, "rb") as comp:
+                fp = io.BytesIO(comp.read())
+        elif compression is None:
+            fp = path
+        reread = read_stata(fp, index_col="index")
+        tm.assert_frame_equal(reread, df)
+
+
+@pytest.mark.parametrize("method", ["zip", "infer"])
+@pytest.mark.parametrize("file_ext", [None, "dta", "zip"])
+def test_compression_dict(method, file_ext):
+    file_name = f"test.{file_ext}"
+    archive_name = "test.dta"
+    df = DataFrame(np.random.randn(10, 2), columns=list("AB"))
+    df.index.name = "index"
+    with tm.ensure_clean(file_name) as path:
+        compression = {"method": method, "archive_name": archive_name}
+        df.to_stata(path, compression=compression)
+        if method == "zip" or file_ext == "zip":
+            zp = zipfile.ZipFile(path, "r")
+            assert len(zp.filelist) == 1
+            assert zp.filelist[0].filename == archive_name
+            fp = io.BytesIO(zp.read(zp.filelist[0]))
+        else:
+            fp = path
+        reread = read_stata(fp, index_col="index")
+        tm.assert_frame_equal(reread, df)