fix ruff errors

Author: kosiew

python/datafusion/__init__.py

@@ -26,26 +26,20 @@
 except ImportError:
     import importlib_metadata
 
+# Local module imports
 from . import functions, object_store, substrait
-
-# The following imports are okay to remain as opaque to the user.
 from ._internal import Config
 from .catalog import Catalog, Database, Table
-from .common import (
-    DFSchema,
-)
+from .common import DFSchema
 from .context import (
+    DataframeDisplayConfig,
     RuntimeEnvBuilder,
+    SQLOptions,
     SessionConfig,
     SessionContext,
-    DataframeDisplayConfig,
-    SQLOptions,
 )
 from .dataframe import DataFrame
-from .expr import (
-    Expr,
-    WindowFrame,
-)
+from .expr import Expr, WindowFrame
 from .io import read_avro, read_csv, read_json, read_parquet
 from .plan import ExecutionPlan, LogicalPlan
 from .record_batch import RecordBatch, RecordBatchStream

python/datafusion/context.py

@@ -26,26 +26,25 @@
 except ImportError:
     from typing_extensions import deprecated  # Python 3.12
 
+if TYPE_CHECKING:
+    import pandas as pd
+    import pathlib
+    import polars as pl
+    import pyarrow as pa
+
+    from datafusion.plan import ExecutionPlan, LogicalPlan
+
 from datafusion.catalog import Catalog, Table
 from datafusion.dataframe import DataFrame
 from datafusion.expr import Expr, SortExpr, sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
 from datafusion.udf import AggregateUDF, ScalarUDF, WindowUDF
 
+from ._internal import DataframeDisplayConfig as DataframeDisplayConfigInternal
 from ._internal import RuntimeEnvBuilder as RuntimeEnvBuilderInternal
+from ._internal import SQLOptions as SQLOptionsInternal
 from ._internal import SessionConfig as SessionConfigInternal
 from ._internal import SessionContext as SessionContextInternal
-from ._internal import SQLOptions as SQLOptionsInternal
-from ._internal import DataframeDisplayConfig as DataframeDisplayConfigInternal
-
-if TYPE_CHECKING:
-    import pathlib
-
-    import pandas as pd
-    import polars as pl
-    import pyarrow as pa
-
-    from datafusion.plan import ExecutionPlan, LogicalPlan
 
 
 class ArrowStreamExportable(Protocol):
@@ -131,7 +130,8 @@ def _validate_positive(self, value: int, name: str) -> None:
             ValueError: If the value is not positive
         """
         if value <= 0:
-            raise ValueError(f"{name} must be greater than 0")
+            error_message = f"{name} must be greater than 0"
+            raise ValueError(error_message)
 
     @property
     def max_table_bytes(self) -> int:

python/tests/test_dataframe.py

@@ -29,9 +29,9 @@
     literal,
 )
 from datafusion import functions as f
+from datafusion.context import DataframeDisplayConfig
 from datafusion.expr import Window
 from pyarrow.csv import write_csv
-from datafusion.context import DataframeDisplayConfig
 
 
 @pytest.fixture
@@ -57,6 +57,11 @@ def data():
     return [{"a": 1, "b": "x" * 50, "c": 3}] * 100
 
 
+@pytest.fixture
+def span_expandable_class():
+    return '<span class="expandable" id="'
+
+
 def test_display_config():
     # Test display_config initialization
     config = DataframeDisplayConfig(
@@ -98,7 +103,7 @@ def test_display_config():
         config.max_table_rows_in_repr = -5
 
 
-def test_session_with_display_config(data):
+def test_session_with_display_config(data, span_expandable_class):
     # Test with_display_config returns a new context with updated config
     ctx = SessionContext()
 
@@ -121,19 +126,14 @@ def test_session_with_display_config(data):
     # The HTML representation should be different with different display configs
     assert html_repr != html_repr2
 
-    # Check that the second representation has the short cell data based on the configured length
-    assert f'<span class="expandable" id="' in html_repr2
-    assert f'>{("x" * 10)}</span>' in html_repr2
+    # Check that the second representation has the short cell data based on the
+    # configured length
+    assert span_expandable_class in html_repr2
+    assert f">{('x' * 10)}</span>" in html_repr2
 
 
 def test_display_config_in_init(data):
-    # Test providing display config directly in SessionContext constructor
-    display_config = DataframeDisplayConfig(
-        max_table_bytes=1024,
-        min_table_rows=5,
-        max_cell_length=10,
-        max_table_rows_in_repr=3,
-    )
+    # Test default display config directly in SessionContext constructor
 
     ctx = SessionContext()
     df1 = ctx.from_pylist(data)
@@ -1403,7 +1403,7 @@ def test_display_config_affects_repr(data):
     assert "Data truncated" not in repr_str2
 
 
-def test_display_config_affects_html_repr(data):
+def test_display_config_affects_html_repr(data, span_expandable_class):
     # Create a context with custom display config to show only a small cell length
     ctx = SessionContext().with_display_config(max_cell_length=5)
 
@@ -1415,8 +1415,7 @@ def test_display_config_affects_html_repr(data):
 
     # The cell should be truncated to 5 characters and have expansion button
     assert ">xxxxx" in html_str  # 5 character limit
-    expandable_class = 'class="expandable-container"'
-    assert expandable_class in html_str
+    assert span_expandable_class in html_str
 
     # Create a context with larger cell length limit
     ctx2 = SessionContext().with_display_config(max_cell_length=60)
@@ -1425,7 +1424,7 @@ def test_display_config_affects_html_repr(data):
     html_str2 = df2._repr_html_()
 
     # String shouldn't be truncated (or at least not in the same way)
-    assert expandable_class not in html_str2
+    assert span_expandable_class not in html_str2
 
 
 def test_display_config_rows_limit_in_html(data):