ENH/TST: Allow more keywords to ensure_clean (pandas-dev#30915)

These keywords will be passed through to
tempfile constructor functions.

Follow-up:

pandas-dev#30771

Author: gfyoung

pandas/_testing.py

@@ -473,7 +473,7 @@ def close(fignum=None):
 
 
 @contextmanager
-def ensure_clean(filename=None, return_filelike=False):
+def ensure_clean(filename=None, return_filelike=False, **kwargs):
     """
     Gets a temporary path and agrees to remove on close.
 
@@ -485,23 +485,37 @@ def ensure_clean(filename=None, return_filelike=False):
     return_filelike : bool (default False)
         if True, returns a file-like which is *always* cleaned. Necessary for
         savefig and other functions which want to append extensions.
+    **kwargs
+        Additional keywords passed in for creating a temporary file.
+        :meth:`tempFile.TemporaryFile` is used when `return_filelike` is ``True``.
+        :meth:`tempfile.mkstemp` is used when `return_filelike` is ``False``.
+        Note that the `filename` parameter will be passed in as the `suffix`
+        argument to either function.
+
+    See Also
+    --------
+    tempfile.TemporaryFile
+    tempfile.mkstemp
     """
     filename = filename or ""
     fd = None
 
+    kwargs["suffix"] = filename
+
     if return_filelike:
-        f = tempfile.TemporaryFile(suffix=filename)
+        f = tempfile.TemporaryFile(**kwargs)
+
         try:
             yield f
         finally:
             f.close()
     else:
-        # don't generate tempfile if using a path with directory specified
+        # Don't generate tempfile if using a path with directory specified.
         if len(os.path.dirname(filename)):
             raise ValueError("Can't pass a qualified name to ensure_clean()")
 
         try:
-            fd, filename = tempfile.mkstemp(suffix=filename)
+            fd, filename = tempfile.mkstemp(**kwargs)
         except UnicodeEncodeError:
             import pytest
 

pandas/tests/io/parser/test_common.py

@@ -8,7 +8,6 @@
 from io import StringIO
 import os
 import platform
-from tempfile import TemporaryFile
 from urllib.error import URLError
 
 import numpy as np
@@ -1847,16 +1846,15 @@ def test_temporary_file(all_parsers):
     parser = all_parsers
     data = "0 0"
 
-    new_file = TemporaryFile("w+")
-    new_file.write(data)
-    new_file.flush()
-    new_file.seek(0)
+    with tm.ensure_clean(mode="w+", return_filelike=True) as new_file:
+        new_file.write(data)
+        new_file.flush()
+        new_file.seek(0)
 
-    result = parser.read_csv(new_file, sep=r"\s+", header=None)
-    new_file.close()
+        result = parser.read_csv(new_file, sep=r"\s+", header=None)
 
-    expected = DataFrame([[0, 0]])
-    tm.assert_frame_equal(result, expected)
+        expected = DataFrame([[0, 0]])
+        tm.assert_frame_equal(result, expected)
 
 
 def test_internal_eof_byte(all_parsers):

pandas/tests/io/parser/test_encoding.py

@@ -5,7 +5,6 @@
 
 from io import BytesIO
 import os
-import tempfile
 
 import numpy as np
 import pytest
@@ -164,7 +163,7 @@ def test_encoding_temp_file(all_parsers, utf_value, encoding_fmt, pass_encoding)
 
     expected = DataFrame({"foo": ["bar"]})
 
-    with tempfile.TemporaryFile(mode="w+", encoding=encoding) as f:
+    with tm.ensure_clean(mode="w+", encoding=encoding, return_filelike=True) as f:
         f.write("foo\nbar")
         f.seek(0)