test modification - context manager

skrub/_reporting/_table_report.py

@@ -1,7 +1,6 @@
 import codecs
 import functools
 import json
-import locale
 from pathlib import Path
 
 from ._html import to_html

skrub/_reporting/tests/test_table_report.py

@@ -1,3 +1,4 @@
+import contextlib
 import datetime
 import json
 import re
@@ -128,25 +129,26 @@ def test_duration(df_module):
 
 @pytest.mark.parametrize(
     "filename_type",
-    ["str", "Path", "file_object", "binary_mode"],
+    ["str", "Path", "text_file_object", "binary_file_object"],
 )
 def test_write_html(tmp_path, pd_module, filename_type):
     df = pd_module.make_dataframe({"a": [1, 2], "b": [3, 4]})
     report = TableReport(df)
 
     tmp_file_path = tmp_path / Path("report.html")
 
-    if filename_type == "str":
-        filename = str(tmp_file_path)
-    elif filename_type == "file_object":
-        filename = open(tmp_file_path, "w", encoding="utf-8")
-    elif filename_type == "binary_mode":
-        filename = open(tmp_file_path, "wb")
-    else:
-        filename = tmp_file_path
+    with contextlib.ExitStack() as stack:
+        if filename_type == "str":
+            filename = str(tmp_file_path)
+        elif filename_type == "text_file_object":
+            filename = stack.enter_context(open(tmp_file_path, "w", encoding="utf-8"))
+        elif filename_type == "binary_file_object":
+            filename = stack.enter_context(open(tmp_file_path, "wb"))
+        else:
+            filename = tmp_file_path
 
-    report.write_html(filename)
-    assert tmp_file_path.exists()
+        report.write_html(filename)
+        assert tmp_file_path.exists()
 
     with open(tmp_file_path, "r", encoding="utf-8") as file:
         saved_content = file.read()
@@ -156,17 +158,22 @@ def test_write_html(tmp_path, pd_module, filename_type):
 def test_write_html_with_not_utf8_encoding(tmp_path, pd_module):
     df = pd_module.make_dataframe({"a": [1, 2], "b": [3, 4]})
     report = TableReport(df)
+    tmp_file_path = tmp_path / Path("report.html")
 
-    filename = open(tmp_path / Path("report.html"), "w", encoding="latin-1")
-    encoding = getattr(filename, "encoding", None)
-    with pytest.raises(
-        ValueError,
-        match=(
-            f"If `file` is a text file it should use utf-8 encoding; got: {encoding!r}"
-        ),
-    ):
-        report.write_html(filename)
-        assert not filename.exists()
+    with open(tmp_file_path, "w", encoding="latin-1") as file:
+        encoding = getattr(file, "encoding", None)
+        with pytest.raises(
+            ValueError,
+            match=(
+                "If `file` is a text file it should use utf-8 encoding; got:"
+                f" {encoding!r}"
+            ),
+        ):
+            report.write_html(file)
+
+    with open(tmp_file_path, "r", encoding="latin-1") as file:
+        saved_content = file.read()
+    assert "</html>" not in saved_content
 
 
 def test_verbosity_parameter(df_module, capsys):