Improve handling of literal arguments: closed and interpolation (#…

…4303)
pola-rs · Aug 8, 2022 · e61e569 · e61e569
1 parent 09bb034
commit e61e569
Show file tree

Hide file tree

Showing 9 changed files with 164 additions and 107 deletions.
diff --git a/py-polars/polars/internals/datatypes.py b/py-polars/polars/internals/datatypes.py
@@ -1,7 +1,16 @@
 from __future__ import annotations
 
+import sys
 from typing import Union
 
 from polars import internals as pli
 
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 IntoExpr = Union[int, float, str, "pli.Expr", "pli.Series"]
+
+ClosedWindow = Literal["left", "right", "both", "none"]
+InterpolationMethod = Literal["nearest", "higher", "lower", "midpoint", "linear"]
diff --git a/py-polars/polars/internals/expr.py b/py-polars/polars/internals/expr.py
@@ -5,7 +5,7 @@
 import random
 import sys
 from datetime import date, datetime, timedelta
-from typing import Any, Callable, List, Sequence
+from typing import TYPE_CHECKING, Any, Callable, List, Sequence
 
 from polars import internals as pli
 from polars.datatypes import (
@@ -41,6 +41,9 @@
 else:
     from typing_extensions import Literal
 
+if TYPE_CHECKING:
+    from polars.internals.datatypes import ClosedWindow, InterpolationMethod
+
 
 def selection_to_pyexpr_list(
     exprs: str | Expr | Sequence[str | Expr | pli.Series] | pli.Series,
@@ -2555,19 +2558,21 @@ def is_duplicated(self) -> Expr:
         """
         return wrap_expr(self._pyexpr.is_duplicated())
 
-    def quantile(self, quantile: float, interpolation: str = "nearest") -> Expr:
+    def quantile(
+        self,
+        quantile: float,
+        interpolation: InterpolationMethod = "nearest",
+    ) -> Expr:
         """
         Get quantile value.
 
 
         Parameters
         ----------
         quantile
-            quantile between 0.0 and 1.0
-
-        interpolation
-            interpolation type, options:
-            ['nearest', 'higher', 'lower', 'midpoint', 'linear']
+            Quantile between 0.0 and 1.0.
+        interpolation : {'nearest', 'higher', 'lower', 'midpoint', 'linear'}
+            Interpolation method.
 
         Examples
         --------
@@ -3241,7 +3246,7 @@ def rolling_min(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Apply a rolling min (moving min) over the values in this array.
@@ -3282,9 +3287,8 @@ def rolling_min(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
 
         .. warning::
@@ -3342,7 +3346,7 @@ def rolling_max(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Apply a rolling max (moving max) over the values in this array.
@@ -3383,9 +3387,8 @@ def rolling_max(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change
@@ -3442,7 +3445,7 @@ def rolling_mean(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Apply a rolling mean (moving mean) over the values in this array.
@@ -3483,9 +3486,8 @@ def rolling_mean(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change
@@ -3540,7 +3542,7 @@ def rolling_sum(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Apply a rolling sum (moving sum) over the values in this array.
@@ -3581,9 +3583,8 @@ def rolling_sum(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change
@@ -3640,7 +3641,7 @@ def rolling_std(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Compute a rolling standard deviation.
@@ -3681,9 +3682,8 @@ def rolling_std(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change
@@ -3711,7 +3711,7 @@ def rolling_var(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Compute a rolling variance.
@@ -3752,9 +3752,8 @@ def rolling_var(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change
@@ -3782,7 +3781,7 @@ def rolling_median(
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Compute a rolling median.
@@ -3819,9 +3818,8 @@ def rolling_median(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change
@@ -3845,24 +3843,23 @@ def rolling_median(
     def rolling_quantile(
         self,
         quantile: float,
-        interpolation: str = "nearest",
+        interpolation: InterpolationMethod = "nearest",
         window_size: int | str = 2,
         weights: List[float] | None = None,
         min_periods: int | None = None,
         center: bool = False,
         by: str | None = None,
-        closed: str = "left",
+        closed: ClosedWindow = "left",
     ) -> Expr:
         """
         Compute a rolling quantile.
 
         Parameters
         ----------
         quantile
-            quantile to compute
-        interpolation
-            interpolation type, options:
-            ['nearest', 'higher', 'lower', 'midpoint', 'linear']
+            Quantile between 0.0 and 1.0.
+        interpolation : {'nearest', 'higher', 'lower', 'midpoint', 'linear'}
+            Interpolation method.
         window_size
             The length of the window. Can be a fixed integer size, or a dynamic temporal
             size indicated by the following string language:
@@ -3893,9 +3890,8 @@ def rolling_quantile(
             If the `window_size` is temporal for instance `"5h"` or `"3s`, you must
             set the column that will be used to determine the windows. This column must
             be of dtype `{Date, Datetime}`
-        closed
-            Defines if the temporal window interval is closed or not.
-            Any of {"left", "right", "both" "none"}
+        closed : {'left', 'right', 'both', 'none'}
+            Define whether the temporal window interval is closed or not.
 
         .. warning::
             The dynamic windows functionality is still experimental and may change