Type aliases for string literals (#4355)

py-polars/polars/convert.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
-import sys
-from typing import Any, Mapping, Sequence, overload
+from typing import TYPE_CHECKING, Any, Mapping, Sequence, overload
 
 from polars.internals import DataFrame, Series
 
@@ -26,10 +25,8 @@
 except ImportError:
     _PANDAS_AVAILABLE = False
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
+if TYPE_CHECKING:
+    from polars.internals.type_aliases import Orientation
 
 
 def from_dict(
@@ -115,7 +112,7 @@ def from_dicts(
 def from_records(
     data: Sequence[Sequence[Any]],
     columns: Sequence[str] | None = None,
-    orient: Literal["col", "row"] | None = None,
+    orient: Orientation | None = None,
 ) -> DataFrame:
     """
     Construct a DataFrame from a sequence of sequences.
@@ -129,7 +126,7 @@ def from_records(
     columns : Sequence of str, default None
         Column labels to use for resulting DataFrame. Must match data dimensions.
         If not specified, columns will be named `column_0`, `column_1`, etc.
-    orient : {'col', 'row'}, default None
+    orient : {None, 'col', 'row'}
         Whether to interpret two-dimensional data as columns or as rows. If None,
         the orientation is inferred by matching the columns and data dimensions. If
         this does not yield conclusive results, column orientation is used.
@@ -163,7 +160,7 @@ def from_records(
 def from_numpy(
     data: np.ndarray[Any, Any],
     columns: Sequence[str] | None = None,
-    orient: Literal["col", "row"] | None = None,
+    orient: Orientation | None = None,
 ) -> DataFrame:
     """
     Construct a DataFrame from a numpy ndarray.
@@ -177,7 +174,7 @@ def from_numpy(
     columns : Sequence of str, default None
         Column labels to use for resulting DataFrame. Must match data dimensions.
         If not specified, columns will be named `column_0`, `column_1`, etc.
-    orient : {'col', 'row'}, default None
+    orient : {None, 'col', 'row'}
         Whether to interpret two-dimensional data as columns or as rows. If None,
         the orientation is inferred by matching the columns and data dimensions. If
         this does not yield conclusive results, column orientation is used.

py-polars/polars/internals/__init__.py

@@ -9,7 +9,6 @@
     _scan_ipc_fsspec,
     _scan_parquet_fsspec,
 )
-from polars.internals.datatypes import IntoExpr
 from polars.internals.expr import (
     Expr,
     expr_to_lit_or_expr,
@@ -42,7 +41,6 @@
 __all__ = [
     "DataFrame",
     "Expr",
-    "IntoExpr",
     "LazyFrame",
     "Series",
     "all",

py-polars/polars/internals/construction.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import sys
 from contextlib import suppress
 from datetime import date, datetime, time, timedelta
 from itertools import zip_longest
@@ -30,9 +29,6 @@
 )
 from polars.utils import threadpool_size
 
-if TYPE_CHECKING:
-    import pandas as pd
-
 try:
     from polars.polars import PyDataFrame, PySeries
 
@@ -54,10 +50,10 @@
 except ImportError:
     _PYARROW_AVAILABLE = False
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
+if TYPE_CHECKING:
+    import pandas as pd
+
+    from polars.internals.type_aliases import Orientation
 
 
 ################################
@@ -466,7 +462,7 @@ def dict_to_pydf(
 def sequence_to_pydf(
     data: Sequence[Any],
     columns: ColumnsType | None = None,
-    orient: Literal["col", "row"] | None = None,
+    orient: Orientation | None = None,
 ) -> PyDataFrame:
     """Construct a PyDataFrame from a sequence."""
     data_series: list[PySeries]
@@ -526,7 +522,7 @@ def sequence_to_pydf(
 def numpy_to_pydf(
     data: np.ndarray[Any, Any],
     columns: ColumnsType | None = None,
-    orient: Literal["col", "row"] | None = None,
+    orient: Orientation | None = None,
 ) -> PyDataFrame:
     """Construct a PyDataFrame from a numpy ndarray."""
     shape = data.shape

py-polars/polars/internals/expr.py

@@ -3,7 +3,6 @@
 import copy
 import math
 import random
-import sys
 from datetime import date, datetime, timedelta
 from typing import TYPE_CHECKING, Any, Callable, List, Sequence
 
@@ -35,16 +34,13 @@
 except ImportError:
     _NUMPY_AVAILABLE = False
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
-
 if TYPE_CHECKING:
-    from polars.internals.datatypes import (
+    from polars.internals.type_aliases import (
         ClosedWindow,
-        FillStrategy,
+        FillNullStrategy,
         InterpolationMethod,
+        ToStructStrategy,
+        TransferEncoding,
     )
 
 
@@ -1884,7 +1880,7 @@ def shift_and_fill(
     def fill_null(
         self,
         value: Any | None = None,
-        strategy: FillStrategy | None = None,
+        strategy: FillNullStrategy | None = None,
         limit: int | None = None,
     ) -> Expr:
         """
@@ -5611,7 +5607,7 @@ def tail(self, n: int = 5) -> Expr:
 
     def to_struct(
         self,
-        n_field_strategy: Literal["first_non_null", "max_width"] = "first_non_null",
+        n_field_strategy: ToStructStrategy = "first_non_null",
         name_generator: Callable[[int], str] | None = None,
     ) -> Expr:
         """
@@ -6154,7 +6150,7 @@ def json_path_match(self, json_path: str) -> Expr:
         """
         return wrap_expr(self._pyexpr.str_json_path_match(json_path))
 
-    def decode(self, encoding: Literal["hex", "base64"], strict: bool = False) -> Expr:
+    def decode(self, encoding: TransferEncoding, strict: bool = False) -> Expr:
         """
         Decode a value using the provided encoding.
 
@@ -6195,7 +6191,7 @@ def decode(self, encoding: Literal["hex", "base64"], strict: bool = False) -> Ex
                 f"encoding must be one of {{'hex', 'base64'}}, got {encoding}"
             )
 
-    def encode(self, encoding: Literal["hex", "base64"]) -> Expr:
+    def encode(self, encoding: TransferEncoding) -> Expr:
         """
         Encode a value using the provided encoding.