shrink type hint

pola-rs · Sep 25, 2023 · b31ec3d · b31ec3d
1 parent 4b47d87
commit b31ec3d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git a/py-polars/polars/expr/expr.py b/py-polars/polars/expr/expr.py
@@ -75,13 +75,16 @@
         FillNullStrategy,
         InterpolationMethod,
         IntoExpr,
+        IntoExprColumn,
         MapElementsStrategy,
         NullBehavior,
+        NumericLiteral,
         PolarsDataType,
         PythonLiteral,
         RankMethod,
         RollingInterpolationMethod,
         SearchSortedSide,
+        TemporalLiteral,
         WindowMappingStrategy,
     )
 
@@ -7424,7 +7427,11 @@ def kurtosis(self, *, fisher: bool = True, bias: bool = True) -> Self:
         """
         return self._from_pyexpr(self._pyexpr.kurtosis(fisher, bias))
 
-    def clip(self, lower_bound: IntoExpr, upper_bound: IntoExpr) -> Self:
+    def clip(
+        self,
+        lower_bound: NumericLiteral | TemporalLiteral | IntoExprColumn,
+        upper_bound: NumericLiteral | TemporalLiteral | IntoExprColumn,
+    ) -> Self:
         """
         Clip (limit) the values in an array to a `min` and `max` boundary.
 
@@ -7461,7 +7468,9 @@ def clip(self, lower_bound: IntoExpr, upper_bound: IntoExpr) -> Self:
         upper_bound = parse_as_expression(upper_bound, str_as_lit=True)
         return self._from_pyexpr(self._pyexpr.clip(lower_bound, upper_bound))
 
-    def clip_min(self, lower_bound: IntoExpr) -> Self:
+    def clip_min(
+        self, lower_bound: NumericLiteral | TemporalLiteral | IntoExprColumn
+    ) -> Self:
         """
         Clip (limit) the values in an array to a `min` boundary.
 
@@ -7495,7 +7504,9 @@ def clip_min(self, lower_bound: IntoExpr) -> Self:
         lower_bound = parse_as_expression(lower_bound, str_as_lit=True)
         return self._from_pyexpr(self._pyexpr.clip_min(lower_bound))
 
-    def clip_max(self, upper_bound: IntoExpr) -> Self:
+    def clip_max(
+        self, upper_bound: NumericLiteral | TemporalLiteral | IntoExprColumn
+    ) -> Self:
         """
         Clip (limit) the values in an array to a `max` boundary.
 

diff --git a/py-polars/polars/series/series.py b/py-polars/polars/series/series.py
@@ -121,14 +121,17 @@
         FillNullStrategy,
         InterpolationMethod,
         IntoExpr,
+        IntoExprColumn,
         NullBehavior,
+        NumericLiteral,
         OneOrMoreDataTypes,
         PolarsDataType,
         PythonLiteral,
         RankMethod,
         RollingInterpolationMethod,
         SearchSortedSide,
         SizeUnit,
+        TemporalLiteral,
     )
 
     if sys.version_info >= (3, 11):
@@ -6101,7 +6104,11 @@ def kurtosis(self, *, fisher: bool = True, bias: bool = True) -> float | None:
         """
         return self._s.kurtosis(fisher, bias)
 
-    def clip(self, lower_bound: IntoExpr, upper_bound: IntoExpr) -> Series:
+    def clip(
+        self,
+        lower_bound: NumericLiteral | TemporalLiteral | IntoExprColumn,
+        upper_bound: NumericLiteral | TemporalLiteral | IntoExprColumn,
+    ) -> Series:
         """
         Clip (limit) the values in an array to a `min` and `max` boundary.
 
@@ -6132,7 +6139,9 @@ def clip(self, lower_bound: IntoExpr, upper_bound: IntoExpr) -> Series:
 
         """
 
-    def clip_min(self, lower_bound: IntoExpr) -> Series:
+    def clip_min(
+        self, lower_bound: NumericLiteral | TemporalLiteral | IntoExprColumn
+    ) -> Series:
         """
         Clip (limit) the values in an array to a `min` boundary.
 
@@ -6148,7 +6157,9 @@ def clip_min(self, lower_bound: IntoExpr) -> Series:
 
         """
 
-    def clip_max(self, upper_bound: IntoExpr) -> Series:
+    def clip_max(
+        self, upper_bound: NumericLiteral | TemporalLiteral | IntoExprColumn
+    ) -> Series:
         """
         Clip (limit) the values in an array to a `max` boundary.
 

diff --git a/py-polars/polars/type_aliases.py b/py-polars/polars/type_aliases.py
@@ -60,9 +60,11 @@
 ]
 SchemaDict: TypeAlias = Mapping[str, PolarsDataType]
 
+NumericLiteral: TypeAlias = Union[int, float, Decimal]
+TemporalLiteral: TypeAlias = Union[date, time, datetime, timedelta]
 # Python literal types (can convert into a `lit` expression)
 PythonLiteral: TypeAlias = Union[
-    str, int, float, bool, date, time, datetime, timedelta, bytes, Decimal, List[Any]
+    NumericLiteral, TemporalLiteral, str, bool, bytes, List[Any]
 ]
 # Inputs that can convert into a `col` expression
 IntoExprColumn: TypeAlias = Union["Expr", "Series", str]