feat[python]: Add string literal types for better type checking (#4400)

py-polars/polars/internals/expr.py

@@ -36,6 +36,7 @@
 
 if TYPE_CHECKING:
     from polars.internals.type_aliases import (
+        CategoricalOrdering,
         ClosedWindow,
         EpochTimeUnit,
         FillNullStrategy,
@@ -7129,17 +7130,18 @@ class ExprCatNameSpace:
     def __init__(self, expr: Expr):
         self._pyexpr = expr._pyexpr
 
-    def set_ordering(self, ordering: str) -> Expr:
+    def set_ordering(self, ordering: CategoricalOrdering) -> Expr:
         """
         Determine how this categorical series should be sorted.
 
         Parameters
         ----------
-        ordering
-            One of:
-                - 'physical' -> use the physical representation of the categories to
-                    determine the order (default)
-                - 'lexical' -. use the string values to determine the ordering
+        ordering : {'physical', 'lexical'}
+            Ordering type:
+
+            - 'physical' -> Use the physical representation of the categories to
+                determine the order (default).
+            - 'lexical' -> Use the string values to determine the ordering.
 
         Examples
         --------

py-polars/polars/internals/frame.py

@@ -114,6 +114,7 @@
         FillNullStrategy,
         InterpolationMethod,
         IpcCompression,
+        JoinStrategy,
         NullStrategy,
         Orientation,
         ParallelStrategy,
@@ -3504,7 +3505,7 @@ def join(
         left_on: str | pli.Expr | list[str | pli.Expr] | None = None,
         right_on: str | pli.Expr | list[str | pli.Expr] | None = None,
         on: str | pli.Expr | list[str | pli.Expr] | None = None,
-        how: str = "inner",
+        how: JoinStrategy = "inner",
         suffix: str = "_right",
     ) -> DataFrame:
         """

py-polars/polars/internals/functions.py

@@ -24,7 +24,7 @@
     _DOCUMENTING = True
 
 if TYPE_CHECKING:
-    from polars.internals.type_aliases import ClosedWindow, TimeUnit
+    from polars.internals.type_aliases import ClosedWindow, ConcatMethod, TimeUnit
 
 
 def get_dummies(df: pli.DataFrame) -> pli.DataFrame:
@@ -44,7 +44,7 @@ def get_dummies(df: pli.DataFrame) -> pli.DataFrame:
 def concat(
     items: Sequence[pli.DataFrame],
     rechunk: bool = True,
-    how: str = "vertical",
+    how: ConcatMethod = "vertical",
 ) -> pli.DataFrame:
     ...
 
@@ -53,7 +53,7 @@ def concat(
 def concat(
     items: Sequence[pli.Series],
     rechunk: bool = True,
-    how: str = "vertical",
+    how: ConcatMethod = "vertical",
 ) -> pli.Series:
     ...
 
@@ -62,7 +62,7 @@ def concat(
 def concat(
     items: Sequence[pli.LazyFrame],
     rechunk: bool = True,
-    how: str = "vertical",
+    how: ConcatMethod = "vertical",
 ) -> pli.LazyFrame:
     ...
 
@@ -71,7 +71,7 @@ def concat(
 def concat(
     items: Sequence[pli.Expr],
     rechunk: bool = True,
-    how: str = "vertical",
+    how: ConcatMethod = "vertical",
 ) -> pli.Expr:
     ...
 
@@ -84,7 +84,7 @@ def concat(
         | Sequence[pli.Expr]
     ),
     rechunk: bool = True,
-    how: str = "vertical",
+    how: ConcatMethod = "vertical",
 ) -> pli.DataFrame | pli.Series | pli.LazyFrame | pli.Expr:
     """
     Aggregate multiple Dataframes/Series to a single DataFrame/Series.
@@ -95,11 +95,9 @@ def concat(
         DataFrames/Series/LazyFrames to concatenate.
     rechunk
         rechunk the final DataFrame/Series.
-    how
+    how : {'vertical', 'diagonal', 'horizontal'}
         Only used if the items are DataFrames.
 
-        One of {"vertical", "diagonal", "horizontal"}.
-
         - Vertical: Applies multiple `vstack` operations.
         - Diagonal: Finds a union between the column schemas and fills missing column
             values with null.
@@ -137,7 +135,7 @@ def concat(
             out = pli.wrap_df(_hor_concat_df(items))
         else:
             raise ValueError(
-                f"how should be one of {'vertical', 'diagonal'}, got {how}"
+                f"how must be one of {{'vertical', 'diagonal'}}, got {how}"
             )
     elif isinstance(first, pli.LazyFrame):
         return pli.wrap_ldf(_concat_lf(items, rechunk))

py-polars/polars/internals/lazy_frame.py

@@ -55,6 +55,7 @@
         CsvEncoding,
         FillNullStrategy,
         InterpolationMethod,
+        JoinStrategy,
         ParallelStrategy,
         UniqueKeepStrategy,
     )
@@ -1342,7 +1343,7 @@ def join(
         left_on: str | pli.Expr | list[str | pli.Expr] | None = None,
         right_on: str | pli.Expr | list[str | pli.Expr] | None = None,
         on: str | pli.Expr | list[str | pli.Expr] | None = None,
-        how: str = "inner",
+        how: JoinStrategy = "inner",
         suffix: str = "_right",
         allow_parallel: bool = True,
         force_parallel: bool = False,

py-polars/polars/internals/series.py

@@ -88,6 +88,7 @@
 
 if TYPE_CHECKING:
     from polars.internals.type_aliases import (
+        CategoricalOrdering,
         ComparisonOperator,
         EpochTimeUnit,
         FillNullStrategy,
@@ -5885,17 +5886,18 @@ class CatNameSpace:
     def __init__(self, s: Series):
         self._s = s
 
-    def set_ordering(self, ordering: str) -> Series:
+    def set_ordering(self, ordering: CategoricalOrdering) -> Series:
         """
         Determine how this categorical series should be sorted.
 
         Parameters
         ----------
-        ordering
-            One of:
-                - 'physical' -> use the physical representation of the categories to
-                    determine the order (default)
-                - 'lexical' -. use the string values to determine the ordering
+        ordering : {'physical', 'lexical'}
+            Ordering type:
+
+            - 'physical' -> Use the physical representation of the categories to
+                determine the order (default).
+            - 'lexical' -> Use the string values to determine the ordering.
 
         Examples
         --------

py-polars/polars/internals/type_aliases.py

@@ -20,6 +20,7 @@
 # User-facing string literal types
 # The following all have an equivalent Rust enum with the same name
 AvroCompression: TypeAlias = Literal["uncompressed", "snappy", "deflate"]
+CategoricalOrdering: TypeAlias = Literal["physical", "lexical"]
 ClosedWindow: TypeAlias = Literal["left", "right", "both", "none"]
 CsvEncoding: TypeAlias = Literal["utf8", "utf8-lossy"]
 FillNullStrategy: TypeAlias = Literal[
@@ -44,11 +45,15 @@
 InterpolationMethod: TypeAlias = Literal[
     "nearest", "higher", "lower", "midpoint", "linear"
 ]  # QuantileInterpolOptions
+JoinStrategy: TypeAlias = Literal[
+    "inner", "left", "outer", "semi", "anti", "cross"
+]  # JoinType
 ToStructStrategy: TypeAlias = Literal[
     "first_non_null", "max_width"
 ]  # ListToStructWidthStrategy
 
 # The following have no equivalent on the Rust side
+ConcatMethod = Literal["vertical", "diagonal", "horizontal"]
 EpochTimeUnit = Literal["ns", "us", "ms", "s", "d"]
 Orientation: TypeAlias = Literal["col", "row"]
 TransferEncoding: TypeAlias = Literal["hex", "base64"]

py-polars/tests/test_df.py

@@ -4,7 +4,7 @@
 import typing
 from datetime import datetime, timedelta
 from io import BytesIO
-from typing import Any, Iterator
+from typing import TYPE_CHECKING, Any, Iterator
 
 import numpy as np
 import pyarrow as pa
@@ -13,6 +13,9 @@
 import polars as pl
 from polars.testing import assert_frame_equal, assert_series_equal, columns
 
+if TYPE_CHECKING:
+    from polars.internals.type_aliases import JoinStrategy
+
 
 def test_version() -> None:
     pl.__version__
@@ -690,7 +693,7 @@ def test_concat() -> None:
         _ = pl.concat([])
 
     with pytest.raises(ValueError):
-        pl.concat([df, df], how="rubbish")
+        pl.concat([df, df], how="rubbish")  # type: ignore[call-overload]
 
 
 def test_arg_where() -> None:
@@ -1753,7 +1756,8 @@ def test_join_suffixes() -> None:
     df_a = pl.DataFrame({"A": [1], "B": [1]})
     df_b = pl.DataFrame({"A": [1], "B": [1]})
 
-    for how in ["left", "inner", "outer", "cross"]:
+    join_strategies: list[JoinStrategy] = ["left", "inner", "outer", "cross"]
+    for how in join_strategies:
         # no need for an assert, we error if wrong
         df_a.join(df_b, on="A", suffix="_y", how=how)["B_y"]
 

py-polars/tests/test_joins.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 from datetime import datetime
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pytest
 
 import polars as pl
 
+if TYPE_CHECKING:
+    from polars.internals.type_aliases import JoinStrategy
+
 
 def test_semi_anti_join() -> None:
     df_a = pl.DataFrame({"key": [1, 2, 3], "payload": ["f", "i", None]})
@@ -82,7 +86,8 @@ def test_sorted_merge_joins() -> None:
             df_a = df_a.select(pl.all().reverse())
             df_b = df_b.select(pl.all().reverse())
 
-        for how in ["left", "inner"]:
+        join_strategies: list[JoinStrategy] = ["left", "inner"]
+        for how in join_strategies:
             # hash join
             out_hash_join = df_a.join(df_b, on="a", how=how)
 
@@ -274,7 +279,8 @@ def test_joins_dispatch() -> None:
         [pl.col("date").str.strptime(pl.Date), pl.col("datetime").cast(pl.Datetime)]
     )
 
-    for how in ["left", "inner", "outer"]:
+    join_strategies: list[JoinStrategy] = ["left", "inner", "outer"]
+    for how in join_strategies:
         dfa.join(dfa, on=["a", "b", "date", "datetime"], how=how)
         dfa.join(dfa, on=["date", "datetime"], how=how)
         dfa.join(dfa, on=["date", "datetime", "a"], how=how)