Repository: pandas-dev/pandas

Issue:
Trouble writing to_stata with a GzipFile

#### Problem description

When a Stata dataset writing to a `GzipFile`, the written dataset is all zero/blank.
I think the Pandas would ideally write out the correct information to the `GzipFile` Stata output, or if that's not an easy change, might consider raising an error when the user tries to write to a `GzipFile`.



#### Expected Output

I expected to read back the same data I tried to write, or to get an error when writing.


Here's the table I tried to write to the GzipFile (`df` in the code):


|a | b | c |
|:--|:--|:--|
|1 | 1.5 | "z"|


Here's the table that gets read back (`df_from_gzip` in the code):


|a | b | c |
|:--|:--|:--|
|0 | 0.0 | ""|


I think this is an error in writing, rather than in reading back, because Stata reads the same all-zeros table.



#### Code Sample

```python

import pandas as pd
import gzip
import subprocess


df = pd.DataFrame({
    'a': [1],
    'b': [1.5],
    'c': ["z"]})

# Use GzipFile to write a compressed version:
with gzip.GzipFile("test_gz.dta.gz", mode = "wb") as f:
    df.to_stata(f, write_index = False)

# Use the system gunzip to extract (using GzipFile fails; see attempt below)
subprocess.run(["gunzip", "--keep", "test_gz.dta.gz"])
df_from_gzip = pd.read_stata("test_gz.dta")

print(df)
print(df_from_gzip)

```

#### Other fun facts


- `bz2.BZ2File` and `lzma.LZMAFile` refuse to write dta files, with the error "UnsupportedOperation: Seeking is only supported on files open for reading"
- Everything works for feather files.
- This isn't an issue with `read_stata`; opening the files in Stata itself gives the same results.
- Variable types are retained.
- Value labels for categorical variables *are* written correctly.
- The number of rows is correct, even for larger examples.
- Reading a system-compressed Stata file is fine.

<details>

```python
import bz2
import lzma


# Try to read the compressed file created before -- fails with the message
# "Not a gzipped file (b'\x01\x00')". I'm not sure why, but it's not central
# to this issue.
with gzip.GzipFile("test_gz.dta.gz") as f:
    df2 = pd.read_stata(f)

    
# Writing feather files to these compressed connections works:
with gzip.GzipFile("test_gz.feather.gz", mode = "wb") as f:
    df.to_feather(f)
with bz2.BZ2File("test_bz.feather.bz2", mode = "wb") as f:
    df.to_feather(f)
with lzma.LZMAFile("test_xz.feather.xz", mode = "wb") as f:
    df.to_feather(f)
        

# Next, writing stata files with other compressors fails because the
# file isn't open for reading.
with bz2.BZ2File("test_bz.dta.bz", mode = "wb") as f:
    df.to_stata(f)  # this raises an error
with lzma.LZMAFile("test_xz.dta.xz", mode = "wb") as f:
    df.to_stata(f)  # this also raises an error


# But reading a system-compressed Stata file works:
df.to_stata("test.dta", write_index = False)
subprocess.run(["gzip", "test.dta"])
with gzip.GzipFile("test.dta.gz") as f:
    assert all(pd.read_stata(f) == df)

```

</details>

#### Output of ``pd.show_versions()``

<details>

INSTALLED VERSIONS
------------------
commit: None
python: 3.6.5.final.0
python-bits: 64
OS: Linux
OS-release: 4.15.0-20-generic
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: en_US.UTF-8

pandas: 0.22.0
pytest: None
pip: 9.0.3
setuptools: 39.1.0
Cython: None
numpy: 1.14.3
scipy: 1.1.0
pyarrow: 0.9.0
xarray: None
IPython: 6.4.0
sphinx: None
patsy: 0.5.0
dateutil: 2.7.2
pytz: 2018.4
blosc: None
bottleneck: None
tables: 3.4.3
numexpr: 2.6.5
feather: 0.4.0
matplotlib: 2.2.2
openpyxl: None
xlrd: None
xlwt: None
xlsxwriter: None
lxml: None
bs4: None
html5lib: 1.0.1
sqlalchemy: None
pymysql: None
psycopg2: None
jinja2: 2.10
s3fs: None
fastparquet: None
pandas_gbq: None
pandas_datareader: None
</details>



Discussion:
Thanks! Could you narrow down your example to a minimal example? It's hard to see exactly what the problem is with that long of an input. http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
Sorry about that!  I reorganized things above, hopefully for the better.
Thanks for the update. Agreed something is going on here. Any interest in debugging further? It'll all be in https://github.com/pandas-dev/pandas/blob/master/pandas/io/stata.py probably. I can help narrow it down further if you need.
I can give it a look, but it will take me a bit.  If you (or anyone else reading this) want to get this working sooner, please do!
I can see this is never going to work with the new format Stata dta writer since writing a dta file requires rewriting some values in an area of the file called the map. It might work with the old format since it is pretty linear.  The docstring should be updated to reflect that it one really works with file object, not general file-like objects.
It appears that ndarray.tofile, which is used to write the data, does not work correctly with gzip files.

```
a = np.arange(2**12)
with gzip.GzipFile("test_nb.gz", mode = "wb") as f:
    a.tofile(f)
```

and then

``` 
gunzip test_nb.gz

gzip: test_nb.gz: not in gzip format
```


Ah, that map business is messy.
Unfortunately, BytesIO doesn't work with NumPy tofile either. TempFile is destroyed immediately on close and so neither of these would allow gzipped files to be cleanly written to disk without an intermediate dta file.

Resolving patch:
diff --git a/pandas/core/frame.py b/pandas/core/frame.py
--- a/pandas/core/frame.py
+++ b/pandas/core/frame.py
@@ -1774,8 +1774,11 @@ def to_stata(self, fname, convert_dates=None, write_index=True,
 
         Parameters
         ----------
-        fname : str or buffer
-            String path of file-like object.
+        fname : path (string), buffer or path object
+            string, path object (pathlib.Path or py._path.local.LocalPath) or
+            object implementing a binary write() functions. If using a buffer
+            then the buffer will not be automatically closed after the file
+            data has been written.
         convert_dates : dict
             Dictionary mapping columns containing datetime types to stata
             internal format to use when writing the dates. Options are 'tc',
diff --git a/pandas/io/stata.py b/pandas/io/stata.py
--- a/pandas/io/stata.py
+++ b/pandas/io/stata.py
@@ -1758,11 +1758,25 @@ def value_labels(self):
         return self.value_label_dict
 
 
-def _open_file_binary_write(fname, encoding):
+def _open_file_binary_write(fname):
+    """
+    Open a binary file or no-op if file-like
+
+    Parameters
+    ----------
+    fname : string path, path object or buffer
+
+    Returns
+    -------
+    file : file-like object
+        File object supporting write
+    own : bool
+        True if the file was created, otherwise False
+    """
     if hasattr(fname, 'write'):
         # if 'b' not in fname.mode:
-        return fname
-    return open(fname, "wb")
+        return fname, False
+    return open(fname, "wb"), True
 
 
 def _set_endianness(endianness):
@@ -1899,7 +1913,9 @@ class StataWriter(StataParser):
     ----------
     fname : path (string), buffer or path object
         string, path object (pathlib.Path or py._path.local.LocalPath) or
-        object implementing a binary write() functions.
+        object implementing a binary write() functions. If using a buffer
+        then the buffer will not be automatically closed after the file
+        is written.
 
         .. versionadded:: 0.23.0 support for pathlib, py.path.
 
@@ -1970,6 +1986,7 @@ def __init__(self, fname, data, convert_dates=None, write_index=True,
         self._time_stamp = time_stamp
         self._data_label = data_label
         self._variable_labels = variable_labels
+        self._own_file = True
         # attach nobs, nvars, data, varlist, typlist
         self._prepare_pandas(data)
 
@@ -2183,9 +2200,7 @@ def _prepare_pandas(self, data):
                 self.fmtlist[key] = self._convert_dates[key]
 
     def write_file(self):
-        self._file = _open_file_binary_write(
-            self._fname, self._encoding or self._default_encoding
-        )
+        self._file, self._own_file = _open_file_binary_write(self._fname)
         try:
             self._write_header(time_stamp=self._time_stamp,
                                data_label=self._data_label)
@@ -2205,6 +2220,23 @@ def write_file(self):
             self._write_file_close_tag()
             self._write_map()
         finally:
+            self._close()
+
+    def _close(self):
+        """
+        Close the file if it was created by the writer.
+
+        If a buffer or file-like object was passed in, for example a GzipFile,
+        then leave this file open for the caller to close. In either case,
+        attempt to flush the file contents to ensure they are written to disk
+        (if supported)
+        """
+        # Some file-like objects might not support flush
+        try:
+            self._file.flush()
+        except AttributeError:
+            pass
+        if self._own_file:
             self._file.close()
 
     def _write_map(self):
@@ -2374,7 +2406,7 @@ def _prepare_data(self):
 
     def _write_data(self):
         data = self.data
-        data.tofile(self._file)
+        self._file.write(data.tobytes())
 
     def _null_terminate(self, s, as_string=False):
         null_byte = '\x00'
@@ -2641,7 +2673,9 @@ class StataWriter117(StataWriter):
     ----------
     fname : path (string), buffer or path object
         string, path object (pathlib.Path or py._path.local.LocalPath) or
-        object implementing a binary write() functions.
+        object implementing a binary write() functions. If using a buffer
+        then the buffer will not be automatically closed after the file
+        is written.
     data : DataFrame
         Input to save
     convert_dates : dict
@@ -2879,7 +2913,7 @@ def _write_data(self):
         self._update_map('data')
         data = self.data
         self._file.write(b'<data>')
-        data.tofile(self._file)
+        self._file.write(data.tobytes())
         self._file.write(b'</data>')
 
     def _write_strls(self):


Test patch:
diff --git a/pandas/tests/io/test_stata.py b/pandas/tests/io/test_stata.py
--- a/pandas/tests/io/test_stata.py
+++ b/pandas/tests/io/test_stata.py
@@ -2,6 +2,8 @@
 # pylint: disable=E1101
 
 import datetime as dt
+import io
+import gzip
 import os
 import struct
 import warnings
@@ -1473,3 +1475,28 @@ def test_invalid_date_conversion(self):
             with pytest.raises(ValueError):
                 original.to_stata(path,
                                   convert_dates={'wrong_name': 'tc'})
+
+    @pytest.mark.parametrize('version', [114, 117])
+    def test_nonfile_writing(self, version):
+        # GH 21041
+        bio = io.BytesIO()
+        df = tm.makeDataFrame()
+        df.index.name = 'index'
+        with tm.ensure_clean() as path:
+            df.to_stata(bio, version=version)
+            bio.seek(0)
+            with open(path, 'wb') as dta:
+                dta.write(bio.read())
+            reread = pd.read_stata(path, index_col='index')
+        tm.assert_frame_equal(df, reread)
+
+    def test_gzip_writing(self):
+        # writing version 117 requires seek and cannot be used with gzip
+        df = tm.makeDataFrame()
+        df.index.name = 'index'
+        with tm.ensure_clean() as path:
+            with gzip.GzipFile(path, 'wb') as gz:
+                df.to_stata(gz, version=114)
+            with gzip.GzipFile(path, 'rb') as gz:
+                reread = pd.read_stata(gz, index_col='index')
+        tm.assert_frame_equal(df, reread)