refactor: update SortKey type definition to use union operator and move declaration

Author: kosiew

python/datafusion/expr.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Optional, Sequence
 
 import pyarrow as pa
 
@@ -43,8 +43,6 @@
 # Standard error message for invalid expression types
 EXPR_TYPE_ERROR = "Use col() or lit() to construct expressions"
 
-SortKey = Union["Expr", "SortExpr", str]
-
 # The following are imported from the internal representation. We may choose to
 # give these all proper wrappers, or to simply leave as is. These were added
 # in order to support passing the `test_imports` unit test.
@@ -1340,3 +1338,6 @@ def nulls_first(self) -> bool:
     def __repr__(self) -> str:
         """Generate a string representation of this expression."""
         return self.raw_sort.__repr__()
+
+
+SortKey = Expr | SortExpr | str