refactor: streamline imports and enhance HTML formatter integration in tests

- Removed redundant import of `configure_formatter` in `__init__.py`.
- Added `configure_formatter` to `__all__` in `__init__.py` for better module exposure.
- Cleaned up import statements in `html_formatter.py` for clarity.
- Consolidated import statements in `test_dataframe.py` for improved readability.
- Simplified the `reset_formatter` fixture by removing unnecessary imports and comments.

Author: kosiew

python/datafusion/__init__.py

@@ -45,11 +45,11 @@
     Expr,
     WindowFrame,
 )
+from .html_formatter import configure_formatter
 from .io import read_avro, read_csv, read_json, read_parquet
 from .plan import ExecutionPlan, LogicalPlan
 from .record_batch import RecordBatch, RecordBatchStream
 from .udf import Accumulator, AggregateUDF, ScalarUDF, WindowUDF, udaf, udf, udwf
-from .html_formatter import configure_formatter
 
 __version__ = importlib_metadata.version(__name__)
 
@@ -77,6 +77,7 @@
     "col",
     "column",
     "common",
+    "configure_formatter",
     "expr",
     "functions",
     "lit",
@@ -91,7 +92,6 @@
     "udf",
     "udwf",
     "unparser",
-    "configure_formatter",
 ]
 
 

python/datafusion/html_formatter.py

@@ -1,7 +1,6 @@
 """HTML formatting utilities for DataFusion DataFrames."""
 
-import sys
-from typing import Dict, Optional, Any, Union, List, Callable, Type, Protocol
+from typing import Any, Callable, Dict, List, Optional, Protocol, Type
 
 
 class CellFormatter(Protocol):

python/tests/test_dataframe.py

@@ -28,8 +28,14 @@
     column,
     literal,
 )
-from datafusion import functions as f
+from datafusion import (
+    functions as f,
+)
 from datafusion.expr import Window
+from datafusion.html_formatter import (
+    _default_formatter,
+    configure_formatter,
+)
 from pyarrow.csv import write_csv
 
 
@@ -659,28 +665,18 @@ def test_window_frame_defaults_match_postgres(partitioned_df):
 @pytest.fixture
 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
-
-    # 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):
     """Test configuring the HTML formatter with different options."""
-    from datafusion.html_formatter import configure_formatter
-
     # Configure with custom settings
     configure_formatter(
         max_cell_length=5,
@@ -700,7 +696,6 @@ def test_html_formatter_configuration(df, reset_formatter):
 
 def test_html_formatter_custom_style_provider(df, reset_formatter):
     """Test using custom style providers with the HTML formatter."""
-    from datafusion.html_formatter import configure_formatter, StyleProvider
 
     class CustomStyleProvider:
         def get_cell_style(self) -> str:
@@ -753,7 +748,7 @@ def custom_cell_builder(value, row, col, table_id):
             num_value = int(value)
             if num_value > 5:  # Values > 5 get green background
                 return f'<td style="background-color: #d9f0d3">{value}</td>'
-            elif num_value < 3:  # Values < 3 get light blue background
+            if num_value < 3:  # Values < 3 get light blue background
                 return f'<td style="background-color: #d3e9f0">{value}</td>'
         except (ValueError, TypeError):
             pass
@@ -804,7 +799,6 @@ def test_html_formatter_complex_customization(df, reset_formatter):
     """Test combining multiple customization options together."""
     from datafusion.html_formatter import (
         configure_formatter,
-        StyleProvider,
         get_formatter,
     )