refactor: update type hints to use typing.Union for better clarity and consistency

Author: kosiew

python/datafusion/expr.py

@@ -22,7 +22,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Optional
+import typing as _typing
+from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Optional, Sequence
 
 import pyarrow as pa
 
@@ -227,7 +228,7 @@
 ]
 
 
-def ensure_expr(value: Expr | Any) -> expr_internal.Expr:
+def ensure_expr(value: _typing.Union["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
@@ -249,7 +250,7 @@ def ensure_expr(value: Expr | Any) -> expr_internal.Expr:
 
 
 def ensure_expr_list(
-    exprs: Iterable[Expr | Iterable[Expr]],
+    exprs: Iterable[_typing.Union["Expr", Iterable["Expr"]]],
 ) -> list[expr_internal.Expr]:
     """Flatten an iterable of expressions, validating each via ``ensure_expr``.
 
@@ -263,7 +264,7 @@ def ensure_expr_list(
         TypeError: If any item is not an instance of :class:`Expr`.
     """
 
-    def _iter(items: Iterable[Expr | Iterable[Expr]]) -> Iterable[expr_internal.Expr]:
+    def _iter(items: Iterable[_typing.Union["Expr", Iterable["Expr"]]]) -> Iterable[expr_internal.Expr]:
         for expr in items:
             if isinstance(expr, Iterable) and not isinstance(
                 expr, (Expr, str, bytes, bytearray)
@@ -276,7 +277,7 @@ def _iter(items: Iterable[Expr | Iterable[Expr]]) -> Iterable[expr_internal.Expr
     return list(_iter(exprs))
 
 
-def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
+def _to_raw_expr(value: _typing.Union["Expr", str]) -> expr_internal.Expr:
     """Convert a Python expression or column name to its raw variant.
 
     Args:
@@ -300,7 +301,7 @@ def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
 
 
 def expr_list_to_raw_expr_list(
-    expr_list: Optional[Sequence[Expr | str] | Expr | str],
+    expr_list: Optional[_typing.Union[Sequence[_typing.Union["Expr", str]], "Expr", str]],
 ) -> Optional[list[expr_internal.Expr]]:
     """Convert a sequence of expressions or column names to raw expressions."""
     if isinstance(expr_list, (Expr, str)):
@@ -310,15 +311,15 @@ def expr_list_to_raw_expr_list(
     return [_to_raw_expr(e) for e in expr_list]
 
 
-def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
+def sort_or_default(e: _typing.Union["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[Sequence[SortKey] | SortKey],
+    sort_list: Optional[_typing.Union[Sequence["SortKey"], "SortKey"]],
 ) -> Optional[list[expr_internal.SortExpr]]:
     """Helper function to return an optional sort list to raw variant."""
     if isinstance(sort_list, (Expr, SortExpr, str)):
@@ -447,7 +448,7 @@ def __invert__(self) -> Expr:
         """Binary not (~)."""
         return Expr(self.expr.__invert__())
 
-    def __getitem__(self, key: str | int) -> Expr:
+    def __getitem__(self, key: _typing.Union[str, int]) -> "Expr":
         """Retrieve sub-object.
 
         If ``key`` is a string, returns the subfield of the struct.
@@ -598,13 +599,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: Any | Expr | None = None) -> Expr:
+    def fill_nan(self, value: Optional[_typing.Union[Any, "Expr"]] = 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: Any | Expr | None = None) -> Expr:
+    def fill_null(self, value: Optional[_typing.Union[Any, "Expr"]] = None) -> "Expr":
         """Fill NULL values with a provided value."""
         if not isinstance(value, Expr):
             value = Expr.literal(value)
@@ -617,7 +618,7 @@ def fill_null(self, value: Any | Expr | None = None) -> Expr:
         bool: pa.bool_(),
     }
 
-    def cast(self, to: pa.DataType[Any] | type[float | int | str | bool]) -> Expr:
+    def cast(self, to: _typing.Union[pa.DataType[Any], type]) -> "Expr":
         """Cast to a new data type."""
         if not isinstance(to, pa.DataType):
             try:
@@ -690,7 +691,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: Expr | SortExpr) -> ExprFuncBuilder:
+    def order_by(self, *exprs: _typing.Union["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
@@ -1239,9 +1240,9 @@ class Window:
 
     def __init__(
         self,
-        partition_by: Optional[list[Expr] | Expr] = None,
+        partition_by: Optional[_typing.Union[list["Expr"], "Expr"]] = None,
         window_frame: Optional[WindowFrame] = None,
-        order_by: Optional[list[SortExpr | Expr | str] | Expr | SortExpr | str] = None,
+        order_by: Optional[_typing.Union[list[_typing.Union["SortExpr", "Expr", str]], "Expr", "SortExpr", str]] = None,
         null_treatment: Optional[NullTreatment] = None,
     ) -> None:
         """Construct a window definition.
@@ -1312,7 +1313,7 @@ def __init__(self, frame_bound: expr_internal.WindowFrameBound) -> None:
         """Constructs a window frame bound."""
         self.frame_bound = frame_bound
 
-    def get_offset(self) -> int | None:
+    def get_offset(self) -> Optional[int]:
         """Returns the offset of the window frame."""
         return self.frame_bound.get_offset()
 
@@ -1396,4 +1397,4 @@ def __repr__(self) -> str:
         return self.raw_sort.__repr__()
 
 
-SortKey = Expr | SortExpr | str
+SortKey = _typing.Union[Expr, SortExpr, str]