From 2c9f4ae223913cfcdf5774b3398b8ea1a85a7aa6 Mon Sep 17 00:00:00 2001
From: gfyoung <gfyoung17@gmail.com>
Date: Sat, 11 Jan 2020 02:29:57 -0800
Subject: [PATCH] ENH/TST: Allow more keywords to ensure_clean

These keywords will be passed through to
tempfile constructor functions.

Follow-up:

https://github.com/pandas-dev/pandas/pull/30771
---
 pandas/_testing.py                      | 22 ++++++++++++++++++----
 pandas/tests/io/parser/test_common.py   | 16 +++++++---------
 pandas/tests/io/parser/test_encoding.py |  3 +--
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/pandas/_testing.py b/pandas/_testing.py
index 1fdc5d478aaf6..72ecbe38ce515 100644
--- a/pandas/_testing.py
+++ b/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
 
diff --git a/pandas/tests/io/parser/test_common.py b/pandas/tests/io/parser/test_common.py
index 4c02a37b66455..6c17f40b790ac 100644
--- a/pandas/tests/io/parser/test_common.py
+++ b/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):
diff --git a/pandas/tests/io/parser/test_encoding.py b/pandas/tests/io/parser/test_encoding.py
index 33abf4bb7d9ee..406e7bedfd298 100644
--- a/pandas/tests/io/parser/test_encoding.py
+++ b/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)