Now that we are on Python 3.10 change from union and otional to |

Author: timsaucer

pyproject.toml

@@ -84,9 +84,8 @@ ignore = [
     "FBT002",  # Allow boolean positional args
     "ISC001",  # Recommended to ignore these rules when using with ruff-format
     "SLF001",  # Allow accessing private members
-    "TD002",
+    "TD002",   # Do not require author names in TODO statements
     "TD003",   # Allow TODO lines
-    "UP007",   # Disallowing Union is pedantic
     # TODO: Enable all of the following, but this PR is getting too large already
     "PLR0913",
     "TRY003",

python/datafusion/dataframe.py

@@ -28,8 +28,6 @@
     Any,
     Iterable,
     Literal,
-    Optional,
-    Union,
     overload,
 )
 
@@ -107,7 +105,7 @@ def from_str(cls: type[Compression], value: str) -> Compression:
                 """
             raise ValueError(error_msg) from err
 
-    def get_default_level(self) -> Optional[int]:
+    def get_default_level(self) -> int | None:
         """Get the default compression level for the compression type.
 
         Returns:
@@ -140,24 +138,24 @@ def __init__(
         write_batch_size: int = 1024,
         writer_version: str = "1.0",
         skip_arrow_metadata: bool = False,
-        compression: Optional[str] = "zstd(3)",
-        compression_level: Optional[int] = None,
-        dictionary_enabled: Optional[bool] = True,
+        compression: str | None = "zstd(3)",
+        compression_level: int | None = None,
+        dictionary_enabled: bool | None = True,
         dictionary_page_size_limit: int = 1024 * 1024,
-        statistics_enabled: Optional[str] = "page",
+        statistics_enabled: str | None = "page",
         max_row_group_size: int = 1024 * 1024,
         created_by: str = "datafusion-python",
-        column_index_truncate_length: Optional[int] = 64,
-        statistics_truncate_length: Optional[int] = None,
+        column_index_truncate_length: int | None = 64,
+        statistics_truncate_length: int | None = None,
         data_page_row_count_limit: int = 20_000,
-        encoding: Optional[str] = None,
+        encoding: str | None = None,
         bloom_filter_on_write: bool = False,
-        bloom_filter_fpp: Optional[float] = None,
-        bloom_filter_ndv: Optional[int] = None,
+        bloom_filter_fpp: float | None = None,
+        bloom_filter_ndv: int | None = None,
         allow_single_file_parallelism: bool = True,
         maximum_parallel_row_group_writers: int = 1,
         maximum_buffered_record_batches_per_stream: int = 2,
-        column_specific_options: Optional[dict[str, ParquetColumnOptions]] = None,
+        column_specific_options: dict[str, ParquetColumnOptions] | None = None,
     ) -> None:
         """Initialize the ParquetWriterOptions.
 
@@ -262,13 +260,13 @@ class ParquetColumnOptions:
 
     def __init__(
         self,
-        encoding: Optional[str] = None,
-        dictionary_enabled: Optional[bool] = None,
-        compression: Optional[str] = None,
-        statistics_enabled: Optional[str] = None,
-        bloom_filter_enabled: Optional[bool] = None,
-        bloom_filter_fpp: Optional[float] = None,
-        bloom_filter_ndv: Optional[int] = None,
+        encoding: str | None = None,
+        dictionary_enabled: bool | None = None,
+        compression: str | None = None,
+        statistics_enabled: str | None = None,
+        bloom_filter_enabled: bool | None = None,
+        bloom_filter_fpp: float | None = None,
+        bloom_filter_ndv: int | None = None,
     ) -> None:
         """Initialize the ParquetColumnOptions.
 
@@ -1063,7 +1061,7 @@ def write_parquet(
     def write_parquet(
         self,
         path: str | pathlib.Path,
-        compression: Union[str, Compression, ParquetWriterOptions] = Compression.ZSTD,
+        compression: str | Compression | ParquetWriterOptions = Compression.ZSTD,
         compression_level: int | None = None,
         write_options: DataFrameWriteOptions | None = None,
     ) -> None:

python/datafusion/dataframe_formatter.py

@@ -21,7 +21,6 @@
 from typing import (
     Any,
     Callable,
-    Optional,
     Protocol,
     runtime_checkable,
 )
@@ -144,9 +143,9 @@ def __init__(
         min_rows_display: int = 20,
         repr_rows: int = 10,
         enable_cell_expansion: bool = True,
-        custom_css: Optional[str] = None,
+        custom_css: str | None = None,
         show_truncation_message: bool = True,
-        style_provider: Optional[StyleProvider] = None,
+        style_provider: StyleProvider | None = None,
         use_shared_styles: bool = True,
     ) -> None:
         """Initialize the HTML formatter.
@@ -226,8 +225,8 @@ def __init__(
         # Registry for custom type formatters
         self._type_formatters: dict[type, CellFormatter] = {}
         # Custom cell builders
-        self._custom_cell_builder: Optional[Callable[[Any, int, int, str], str]] = None
-        self._custom_header_builder: Optional[Callable[[Any], str]] = None
+        self._custom_cell_builder: Callable[[Any, int, int, str], str] | None = None
+        self._custom_header_builder: Callable[[Any], str] | None = None
 
     def register_formatter(self, type_class: type, formatter: CellFormatter) -> None:
         """Register a custom formatter for a specific data type.

python/datafusion/expr.py

@@ -22,8 +22,7 @@
 
 from __future__ import annotations
 
-import typing as _typing
-from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Optional, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Sequence
 
 try:
     from warnings import deprecated  # Python 3.13+
@@ -230,7 +229,7 @@
 ]
 
 
-def ensure_expr(value: _typing.Union[Expr, Any]) -> expr_internal.Expr:
+def ensure_expr(value: Expr | Any) -> expr_internal.Expr:
     """Return the internal expression from ``Expr`` or raise ``TypeError``.
 
     This helper rejects plain strings and other non-:class:`Expr` values so
@@ -252,7 +251,7 @@ def ensure_expr(value: _typing.Union[Expr, Any]) -> expr_internal.Expr:
 
 
 def ensure_expr_list(
-    exprs: Iterable[_typing.Union[Expr, Iterable[Expr]]],
+    exprs: Iterable[Expr | Iterable[Expr]],
 ) -> list[expr_internal.Expr]:
     """Flatten an iterable of expressions, validating each via ``ensure_expr``.
 
@@ -267,7 +266,7 @@ def ensure_expr_list(
     """
 
     def _iter(
-        items: Iterable[_typing.Union[Expr, Iterable[Expr]]],
+        items: Iterable[Expr | Iterable[Expr]],
     ) -> Iterable[expr_internal.Expr]:
         for expr in items:
             if isinstance(expr, Iterable) and not isinstance(
@@ -281,7 +280,7 @@ def _iter(
     return list(_iter(exprs))
 
 
-def _to_raw_expr(value: _typing.Union[Expr, str]) -> expr_internal.Expr:
+def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
     """Convert a Python expression or column name to its raw variant.
 
     Args:
@@ -305,8 +304,8 @@ def _to_raw_expr(value: _typing.Union[Expr, str]) -> expr_internal.Expr:
 
 
 def expr_list_to_raw_expr_list(
-    expr_list: Optional[list[Expr] | Expr],
-) -> Optional[list[expr_internal.Expr]]:
+    expr_list: list[Expr] | Expr | None,
+) -> list[expr_internal.Expr] | None:
     """Convert a sequence of expressions or column names to raw expressions."""
     if isinstance(expr_list, Expr | str):
         expr_list = [expr_list]
@@ -315,16 +314,16 @@ def expr_list_to_raw_expr_list(
     return [_to_raw_expr(e) for e in expr_list]
 
 
-def sort_or_default(e: _typing.Union[Expr, SortExpr]) -> expr_internal.SortExpr:
+def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
     """Helper function to return a default Sort if an Expr is provided."""
     if isinstance(e, SortExpr):
         return e.raw_sort
     return SortExpr(e, ascending=True, nulls_first=True).raw_sort
 
 
 def sort_list_to_raw_sort_list(
-    sort_list: Optional[_typing.Union[Sequence[SortKey], SortKey]],
-) -> Optional[list[expr_internal.SortExpr]]:
+    sort_list: Sequence[SortKey] | SortKey | None,
+) -> list[expr_internal.SortExpr] | None:
     """Helper function to return an optional sort list to raw variant."""
     if isinstance(sort_list, Expr | SortExpr | str):
         sort_list = [sort_list]
@@ -601,7 +600,7 @@ def column(value: str) -> Expr:
         """Creates a new expression representing a column."""
         return Expr(expr_internal.RawExpr.column(value))
 
-    def alias(self, name: str, metadata: Optional[dict[str, str]] = None) -> Expr:
+    def alias(self, name: str, metadata: dict[str, str] | None = None) -> Expr:
         """Assign a name to the expression.
 
         Args:
@@ -630,13 +629,13 @@ def is_not_null(self) -> Expr:
         """Returns ``True`` if this expression is not null."""
         return Expr(self.expr.is_not_null())
 
-    def fill_nan(self, value: Optional[_typing.Union[Any, Expr]] = None) -> Expr:
+    def fill_nan(self, value: Any | Expr | None = None) -> Expr:
         """Fill NaN values with a provided value."""
         if not isinstance(value, Expr):
             value = Expr.literal(value)
         return Expr(functions_internal.nanvl(self.expr, value.expr))
 
-    def fill_null(self, value: Optional[_typing.Union[Any, Expr]] = None) -> Expr:
+    def fill_null(self, value: Any | Expr | None = None) -> Expr:
         """Fill NULL values with a provided value."""
         if not isinstance(value, Expr):
             value = Expr.literal(value)
@@ -649,7 +648,7 @@ def fill_null(self, value: Optional[_typing.Union[Any, Expr]] = None) -> Expr:
         bool: pa.bool_(),
     }
 
-    def cast(self, to: _typing.Union[pa.DataType[Any], type]) -> Expr:
+    def cast(self, to: pa.DataType[Any] | type) -> Expr:
         """Cast to a new data type."""
         if not isinstance(to, pa.DataType):
             try:
@@ -722,7 +721,7 @@ def column_name(self, plan: LogicalPlan) -> str:
         """Compute the output column name based on the provided logical plan."""
         return self.expr.column_name(plan._raw_plan)
 
-    def order_by(self, *exprs: _typing.Union[Expr, SortExpr]) -> ExprFuncBuilder:
+    def order_by(self, *exprs: Expr | SortExpr) -> ExprFuncBuilder:
         """Set the ordering for a window or aggregate function.
 
         This function will create an :py:class:`ExprFuncBuilder` that can be used to
@@ -1271,17 +1270,10 @@ class Window:
 
     def __init__(
         self,
-        partition_by: Optional[_typing.Union[list[Expr], Expr]] = None,
-        window_frame: Optional[WindowFrame] = None,
-        order_by: Optional[
-            _typing.Union[
-                list[_typing.Union[SortExpr, Expr, str]],
-                Expr,
-                SortExpr,
-                str,
-            ]
-        ] = None,
-        null_treatment: Optional[NullTreatment] = None,
+        partition_by: list[Expr] | Expr | None = None,
+        window_frame: WindowFrame | None = None,
+        order_by: list[SortExpr | Expr | str] | Expr | SortExpr | str | None = None,
+        null_treatment: NullTreatment | None = None,
     ) -> None:
         """Construct a window definition.
 
@@ -1301,7 +1293,7 @@ class WindowFrame:
     """Defines a window frame for performing window operations."""
 
     def __init__(
-        self, units: str, start_bound: Optional[Any], end_bound: Optional[Any]
+        self, units: str, start_bound: Any | None, end_bound: Any | None
     ) -> None:
         """Construct a window frame using the given parameters.
 
@@ -1351,7 +1343,7 @@ def __init__(self, frame_bound: expr_internal.WindowFrameBound) -> None:
         """Constructs a window frame bound."""
         self.frame_bound = frame_bound
 
-    def get_offset(self) -> Optional[int]:
+    def get_offset(self) -> int | None:
         """Returns the offset of the window frame."""
         return self.frame_bound.get_offset()
 
@@ -1435,4 +1427,4 @@ def __repr__(self) -> str:
         return self.raw_sort.__repr__()
 
 
-SortKey = _typing.Union[Expr, SortExpr, str]
+SortKey = Expr | SortExpr | str