BUG: Ensure incomplete stata files are deleted

Attempt to delete failed writes and warn if not able to delete

Author: bashtage

doc/source/whatsnew/v0.24.0.rst

@@ -1576,6 +1576,7 @@ Notice how we now instead output ``np.nan`` itself instead of a stringified form
 - :func:`DataFrame.to_string()`, :func:`DataFrame.to_html()`, :func:`DataFrame.to_latex()` will correctly format output when a string is passed as the ``float_format`` argument (:issue:`21625`, :issue:`22270`)
 - Bug in :func:`read_csv` that caused it to raise ``OverflowError`` when trying to use 'inf' as ``na_value`` with integer index column (:issue:`17128`)
 - Bug in :func:`pandas.io.json.json_normalize` that caused it to raise ``TypeError`` when two consecutive elements of ``record_path`` are dicts (:issue:`22706`)
+- Bug in :meth:`DataFrame.to_stata`, :class:`pandas.io.stata.StataWriter` and :class:`pandas.io.stata.StataWriter117` where a exception would leave a partially written and invalid dta file (:issue:`23573`)
 - Bug in :meth:`DataFrame.to_stata` and :class:`pandas.io.stata.StataWriter117` that produced invalid files when using strLs with non-ASCII characters (:issue:`23573`)
 
 Plotting

pandas/io/stata.py

@@ -12,6 +12,7 @@
 
 from collections import OrderedDict
 import datetime
+import os
 import struct
 import sys
 import warnings
@@ -503,6 +504,10 @@ class InvalidColumnName(Warning):
 """
 
 
+class IOWarning(Warning):
+    pass
+
+
 def _cast_to_stata_types(data):
     """Checks the dtypes of the columns of a pandas DataFrame for
     compatibility with the data types and ranges supported by Stata, and
@@ -2209,7 +2214,17 @@ def write_file(self):
             self._write_value_labels()
             self._write_file_close_tag()
             self._write_map()
-        finally:
+        except Exception as exc:
+            self._close()
+            try:
+                if self._own_file:
+                    os.unlink(self._fname)
+            except Exception:
+                warnings.warn('This save was not successful but {0} could not '
+                              'be deleted.  This file is not '
+                              'valid.'.format(self._fname), IOWarning)
+            raise exc
+        else:
             self._close()
 
     def _close(self):

pandas/tests/io/test_stata.py

@@ -21,7 +21,8 @@
 from pandas.core.frame import DataFrame, Series
 from pandas.io.parsers import read_csv
 from pandas.io.stata import (InvalidColumnName, PossiblePrecisionLoss,
-                             StataMissingValue, StataReader, read_stata)
+                             StataMissingValue, StataReader, read_stata,
+                             IOWarning)
 
 
 @pytest.fixture
@@ -1547,6 +1548,16 @@ def test_all_none_exception(self, version):
         assert 'Only string-like' in excinfo.value.args[0]
         assert 'Column `none`' in excinfo.value.args[0]
 
+    @pytest.mark.parametrize('version', [114, 117])
+    def test_invalid_file_not_written(self, version):
+        content = 'Here is one __�__ Another one __·__ Another one __½__'
+        df = DataFrame([content], columns=['invalid'])
+        expected_exc = UnicodeEncodeError if PY3 else UnicodeDecodeError
+        with tm.ensure_clean() as path:
+            with pytest.raises(expected_exc):
+                with tm.assert_produces_warning(IOWarning):
+                    df.to_stata(path)
+
     def test_strl_latin1(self):
         # GH 23573, correct GSO data to reflect correct size
         output = DataFrame([[u'pandas'] * 2, [u'þâÑÐŧ'] * 2],