feat(typing): Add polars literal aliases (#2352)

* feat(typing): Add some `polars` aliases

Towards #2337

* refactor(typing): Replace literals in impls

* feat(typing): Add `LazyUniqueKeepStrategy`

Not from `polars`, but describes the subset we support for `LazyFrame.unique`

* docs(typing): Add trailing docstrings

Copied over unchanged, but seems they could do with standardising

* docs: Update `typing.md`

- Following the existing pattern of `SizeUnit`, `TimeUnit`
- They appear in docs, but **not** in `nw.typing.__all__`

* refactor(typing): Replace literals in public api

* refactor(typing): Update `v1`

* docs: Use more consistent markdown

- Use `-` for bullet lists
- Wrap literals in `"..."`
- Single space after `:`
- Titlecase for sentences
- Inline code for parameter names (e.g. ``on``)
- Wrap lines at ~80 characters

Author: dangotbanned

docs/api-reference/typing.md

@@ -19,6 +19,16 @@ Narwhals comes fully statically typed. In addition to `nw.DataFrame`, `nw.Expr`,
         - IntoSeriesT
         - SizeUnit
         - TimeUnit
+        - AsofJoinStrategy
+        - ClosedInterval
+        - ConcatMethod
+        - FillNullStrategy
+        - JoinStrategy
+        - PivotAgg
+        - RankMethod
+        - RollingInterpolationMethod
+        - UniqueKeepStrategy
+        - LazyUniqueKeepStrategy
       show_source: false
       show_bases: false
 

narwhals/_arrow/dataframe.py

@@ -59,7 +59,9 @@
     from narwhals.schema import Schema
     from narwhals.typing import CompliantDataFrame
     from narwhals.typing import CompliantLazyFrame
+    from narwhals.typing import JoinStrategy
     from narwhals.typing import SizeUnit
+    from narwhals.typing import UniqueKeepStrategy
     from narwhals.typing import _1DArray
     from narwhals.typing import _2DArray
     from narwhals.utils import Version
@@ -451,7 +453,7 @@ def join(
         self: Self,
         other: Self,
         *,
-        how: Literal["inner", "left", "full", "cross", "semi", "anti"],
+        how: JoinStrategy,
         left_on: Sequence[str] | None,
         right_on: Sequence[str] | None,
         suffix: str,
@@ -758,7 +760,7 @@ def unique(
         self: ArrowDataFrame,
         subset: Sequence[str] | None,
         *,
-        keep: Literal["any", "first", "last", "none"],
+        keep: UniqueKeepStrategy,
         maintain_order: bool | None = None,
     ) -> ArrowDataFrame:
         # The param `maintain_order` is only here for compatibility with the Polars API

narwhals/_arrow/expr.py

@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import Literal
 from typing import Sequence
 
 import pyarrow.compute as pc
@@ -25,6 +24,7 @@
     from narwhals._compliant.typing import EvalNames
     from narwhals._compliant.typing import EvalSeries
     from narwhals._expression_parsing import ExprMetadata
+    from narwhals.typing import RankMethod
     from narwhals.utils import Version
     from narwhals.utils import _FullContext
 
@@ -208,12 +208,7 @@ def cum_max(self: Self, *, reverse: bool) -> Self:
     def cum_prod(self: Self, *, reverse: bool) -> Self:
         return self._reuse_series("cum_prod", reverse=reverse)
 
-    def rank(
-        self: Self,
-        method: Literal["average", "min", "max", "dense", "ordinal"],
-        *,
-        descending: bool,
-    ) -> Self:
+    def rank(self, method: RankMethod, *, descending: bool) -> Self:
         return self._reuse_series("rank", method=method, descending=descending)
 
     ewm_mean = not_implemented()

narwhals/_arrow/namespace.py

@@ -34,6 +34,7 @@
     from narwhals._arrow.typing import ArrowChunkedArray
     from narwhals._arrow.typing import Incomplete
     from narwhals.dtypes import DType
+    from narwhals.typing import ConcatMethod
     from narwhals.utils import Version
 
 
@@ -211,10 +212,7 @@ def func(df: ArrowDataFrame) -> list[ArrowSeries]:
         )
 
     def concat(
-        self: Self,
-        items: Iterable[ArrowDataFrame],
-        *,
-        how: Literal["horizontal", "vertical", "diagonal"],
+        self, items: Iterable[ArrowDataFrame], *, how: ConcatMethod
     ) -> ArrowDataFrame:
         dfs = [item.native for item in items]
 

narwhals/_arrow/series.py

@@ -4,7 +4,6 @@
 from typing import Any
 from typing import Iterable
 from typing import Iterator
-from typing import Literal
 from typing import Mapping
 from typing import Sequence
 from typing import cast
@@ -56,7 +55,11 @@
     from narwhals._arrow.typing import _AsPyType
     from narwhals._arrow.typing import _BasicDataType
     from narwhals.dtypes import DType
+    from narwhals.typing import ClosedInterval
+    from narwhals.typing import FillNullStrategy
     from narwhals.typing import Into1DArray
+    from narwhals.typing import RankMethod
+    from narwhals.typing import RollingInterpolationMethod
     from narwhals.typing import _1DArray
     from narwhals.typing import _2DArray
     from narwhals.utils import Version
@@ -499,10 +502,7 @@ def all(self: Self, *, _return_py_scalar: bool = True) -> bool:
         )
 
     def is_between(
-        self: Self,
-        lower_bound: Any,
-        upper_bound: Any,
-        closed: Literal["left", "right", "none", "both"],
+        self, lower_bound: Any, upper_bound: Any, closed: ClosedInterval
     ) -> Self:
         _, lower_bound = extract_native(self, lower_bound)
         _, upper_bound = extract_native(self, upper_bound)
@@ -636,17 +636,14 @@ def sample(
         return self._with_native(self.native.take(mask))
 
     def fill_null(
-        self: Self,
-        value: Any | None,
-        strategy: Literal["forward", "backward"] | None,
-        limit: int | None,
+        self, value: Any | None, strategy: FillNullStrategy | None, limit: int | None
     ) -> Self:
         import numpy as np  # ignore-banned-import
 
         def fill_aux(
             arr: ArrowArray | ArrowChunkedArray,
             limit: int,
-            direction: Literal["forward", "backward"] | None = None,
+            direction: FillNullStrategy | None = None,
         ) -> ArrowArray:
             # this algorithm first finds the indices of the valid values to fill all the null value positions
             # then it calculates the distance of each new index and the original index
@@ -812,9 +809,9 @@ def to_dummies(self: Self, *, separator: str, drop_first: bool) -> ArrowDataFram
         ).simple_select(*output_order)
 
     def quantile(
-        self: Self,
+        self,
         quantile: float,
-        interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"],
+        interpolation: RollingInterpolationMethod,
         *,
         _return_py_scalar: bool = True,
     ) -> float:
@@ -1028,12 +1025,7 @@ def rolling_std(
             ** 0.5
         )
 
-    def rank(
-        self: Self,
-        method: Literal["average", "min", "max", "dense", "ordinal"],
-        *,
-        descending: bool,
-    ) -> Self:
+    def rank(self, method: RankMethod, *, descending: bool) -> Self:
         if method == "average":
             msg = (
                 "`rank` with `method='average' is not supported for pyarrow backend. "

narwhals/_compliant/dataframe.py

@@ -40,7 +40,11 @@
     from narwhals._translate import IntoArrowTable
     from narwhals.dtypes import DType
     from narwhals.schema import Schema
+    from narwhals.typing import AsofJoinStrategy
+    from narwhals.typing import JoinStrategy
+    from narwhals.typing import LazyUniqueKeepStrategy
     from narwhals.typing import SizeUnit
+    from narwhals.typing import UniqueKeepStrategy
     from narwhals.typing import _2DArray
     from narwhals.utils import Implementation
     from narwhals.utils import _FullContext
@@ -144,7 +148,7 @@ def join(
         self: Self,
         other: Self,
         *,
-        how: Literal["inner", "left", "full", "cross", "semi", "anti"],
+        how: JoinStrategy,
         left_on: Sequence[str] | None,
         right_on: Sequence[str] | None,
         suffix: str,
@@ -157,7 +161,7 @@ def join_asof(
         right_on: str | None,
         by_left: Sequence[str] | None,
         by_right: Sequence[str] | None,
-        strategy: Literal["backward", "forward", "nearest"],
+        strategy: AsofJoinStrategy,
         suffix: str,
     ) -> Self: ...
     def lazy(self, *, backend: Implementation | None) -> CompliantLazyFrame[Any, Any]: ...
@@ -193,7 +197,7 @@ def unique(
         self,
         subset: Sequence[str] | None,
         *,
-        keep: Literal["any", "first", "last", "none"],
+        keep: UniqueKeepStrategy,
         maintain_order: bool | None = None,
     ) -> Self: ...
     def unpivot(
@@ -284,7 +288,7 @@ def join_asof(
         right_on: str | None,
         by_left: Sequence[str] | None,
         by_right: Sequence[str] | None,
-        strategy: Literal["backward", "forward", "nearest"],
+        strategy: AsofJoinStrategy,
         suffix: str,
     ) -> Self: ...
     def rename(self, mapping: Mapping[str, str]) -> Self: ...
@@ -295,10 +299,7 @@ def sort(
     @deprecated("`LazyFrame.tail` is deprecated and will be removed in a future version.")
     def tail(self, n: int) -> Self: ...
     def unique(
-        self,
-        subset: Sequence[str] | None,
-        *,
-        keep: Literal["any", "none"],
+        self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
     ) -> Self: ...
     def unpivot(
         self,

narwhals/_compliant/expr.py

@@ -64,6 +64,9 @@
     from narwhals._expression_parsing import ExprKind
     from narwhals._expression_parsing import ExprMetadata
     from narwhals.dtypes import DType
+    from narwhals.typing import FillNullStrategy
+    from narwhals.typing import RankMethod
+    from narwhals.typing import RollingInterpolationMethod
     from narwhals.typing import TimeUnit
     from narwhals.utils import Implementation
     from narwhals.utils import Version
@@ -131,10 +134,7 @@ def n_unique(self) -> Self: ...
     def null_count(self) -> Self: ...
     def drop_nulls(self) -> Self: ...
     def fill_null(
-        self,
-        value: Any | None,
-        strategy: Literal["forward", "backward"] | None,
-        limit: int | None,
+        self, value: Any | None, strategy: FillNullStrategy | None, limit: int | None
     ) -> Self: ...
     def diff(self) -> Self: ...
     def unique(self) -> Self: ...
@@ -156,12 +156,7 @@ def cum_max(self, *, reverse: bool) -> Self: ...
     def cum_prod(self, *, reverse: bool) -> Self: ...
     def is_in(self, other: Any) -> Self: ...
     def sort(self, *, descending: bool, nulls_last: bool) -> Self: ...
-    def rank(
-        self,
-        method: Literal["average", "min", "max", "dense", "ordinal"],
-        *,
-        descending: bool,
-    ) -> Self: ...
+    def rank(self, method: RankMethod, *, descending: bool) -> Self: ...
     def replace_strict(
         self,
         old: Sequence[Any] | Mapping[Any, Any],
@@ -181,9 +176,7 @@ def sample(
         seed: int | None,
     ) -> Self: ...
     def quantile(
-        self,
-        quantile: float,
-        interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"],
+        self, quantile: float, interpolation: RollingInterpolationMethod
     ) -> Self: ...
     def map_batches(
         self,
@@ -657,10 +650,7 @@ def is_nan(self) -> Self:
         return self._reuse_series("is_nan")
 
     def fill_null(
-        self,
-        value: Any | None,
-        strategy: Literal["forward", "backward"] | None,
-        limit: int | None,
+        self, value: Any | None, strategy: FillNullStrategy | None, limit: int | None
     ) -> Self:
         return self._reuse_series(
             "fill_null", value=value, strategy=strategy, limit=limit
@@ -746,9 +736,7 @@ def is_last_distinct(self) -> Self:
         return self._reuse_series("is_last_distinct")
 
     def quantile(
-        self,
-        quantile: float,
-        interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"],
+        self, quantile: float, interpolation: RollingInterpolationMethod
     ) -> Self:
         return self._reuse_series(
             "quantile",

narwhals/_compliant/namespace.py

@@ -5,7 +5,6 @@
 from typing import Any
 from typing import Container
 from typing import Iterable
-from typing import Literal
 from typing import Mapping
 from typing import Protocol
 from typing import Sequence
@@ -35,6 +34,7 @@
     from narwhals._compliant.when_then import EagerWhen
     from narwhals.dtypes import DType
     from narwhals.schema import Schema
+    from narwhals.typing import ConcatMethod
     from narwhals.typing import Into1DArray
     from narwhals.typing import _2DArray
     from narwhals.utils import Implementation
@@ -75,10 +75,7 @@ def mean_horizontal(self, *exprs: CompliantExprT) -> CompliantExprT: ...
     def min_horizontal(self, *exprs: CompliantExprT) -> CompliantExprT: ...
     def max_horizontal(self, *exprs: CompliantExprT) -> CompliantExprT: ...
     def concat(
-        self,
-        items: Iterable[CompliantFrameT],
-        *,
-        how: Literal["horizontal", "vertical", "diagonal"],
+        self, items: Iterable[CompliantFrameT], *, how: ConcatMethod
     ) -> CompliantFrameT: ...
     def when(
         self, predicate: CompliantExprT

narwhals/_compliant/series.py

@@ -5,7 +5,6 @@
 from typing import Generic
 from typing import Iterable
 from typing import Iterator
-from typing import Literal
 from typing import Mapping
 from typing import Protocol
 from typing import Sequence
@@ -40,7 +39,11 @@
     from narwhals._compliant.namespace import CompliantNamespace
     from narwhals._compliant.namespace import EagerNamespace
     from narwhals.dtypes import DType
+    from narwhals.typing import ClosedInterval
+    from narwhals.typing import FillNullStrategy
     from narwhals.typing import Into1DArray
+    from narwhals.typing import RankMethod
+    from narwhals.typing import RollingInterpolationMethod
     from narwhals.typing import _1DArray
     from narwhals.utils import Implementation
     from narwhals.utils import Version
@@ -152,10 +155,7 @@ def ewm_mean(
         ignore_nulls: bool,
     ) -> Self: ...
     def fill_null(
-        self,
-        value: Any | None,
-        strategy: Literal["forward", "backward"] | None,
-        limit: int | None,
+        self, value: Any | None, strategy: FillNullStrategy | None, limit: int | None
     ) -> Self: ...
     def filter(self, predicate: Any) -> Self: ...
     def gather_every(self, n: int, offset: int) -> Self: ...
@@ -169,10 +169,7 @@ def hist(
     ) -> CompliantDataFrame[Self, Any, Any]: ...
     def head(self, n: int) -> Self: ...
     def is_between(
-        self,
-        lower_bound: Any,
-        upper_bound: Any,
-        closed: Literal["left", "right", "none", "both"],
+        self, lower_bound: Any, upper_bound: Any, closed: ClosedInterval
     ) -> Self: ...
     def is_finite(self) -> Self: ...
     def is_first_distinct(self) -> Self: ...
@@ -192,16 +189,9 @@ def mode(self) -> Self: ...
     def n_unique(self) -> int: ...
     def null_count(self) -> int: ...
     def quantile(
-        self,
-        quantile: float,
-        interpolation: Literal["nearest", "higher", "lower", "midpoint", "linear"],
+        self, quantile: float, interpolation: RollingInterpolationMethod
     ) -> float: ...
-    def rank(
-        self,
-        method: Literal["average", "min", "max", "dense", "ordinal"],
-        *,
-        descending: bool,
-    ) -> Self: ...
+    def rank(self, method: RankMethod, *, descending: bool) -> Self: ...
     def replace_strict(
         self,
         old: Sequence[Any] | Mapping[Any, Any],