refactor: improve HTML formatting logic in DataFrame by separating data collection and schema retrieval for clarity

kosiew · kosiew · commit 2c3bd604e348 · 2025-04-08T12:16:56.000+08:00
refactor: enhance reset_formatter fixture to preserve original formatter configuration during tests
diff --git a/python/datafusion/dataframe.py b/python/datafusion/dataframe.py
@@ -156,11 +156,15 @@ def _repr_html_(self) -> str:
         # Import here to avoid circular imports
         from datafusion.html_formatter import get_formatter
 
-        # Always get the latest formatter
+        # Always get the latest formatter instance
         formatter = get_formatter()
 
-        # Format the data using the latest formatter
-        return formatter.format_html(self.collect(), self.schema())
+        # Get data and schema
+        batches = self.collect()
+        schema = self.schema()
+
+        # Format the data using our formatter
+        return formatter.format_html(batches, schema)
 
     def describe(self) -> DataFrame:
         """Return the statistics for this DataFrame.
diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py
@@ -661,8 +661,20 @@ def reset_formatter():
     """Reset the HTML formatter after each test."""
     from datafusion.html_formatter import configure_formatter
 
+    # Store original formatter configuration
+    from datafusion.html_formatter import _default_formatter
+
+    original = _default_formatter
+
+    # Give the test a fresh formatter
+    configure_formatter()
+
     yield
-    configure_formatter()  # Reset to defaults after test
+
+    # Completely reset to original state after test
+    from datafusion.html_formatter import _default_formatter
+
+    globals()["_default_formatter"] = original
 
 
 def test_html_formatter_configuration(df, reset_formatter):
