Turn on disallow-any-generics (#4320)

Co-authored-by: Matteo Santamaria <msantama@gmail.com>

py-polars/polars/convert.py

@@ -34,7 +34,7 @@
 
 
 def from_dict(
-    data: Mapping[str, Sequence | Mapping],
+    data: Mapping[str, Sequence[object] | Mapping[str, Sequence[object]]],
     columns: Sequence[str] | None = None,
 ) -> DataFrame:
     """
@@ -70,7 +70,7 @@ def from_dict(
     └─────┴─────┘
 
     """
-    return DataFrame._from_dict(data=data, columns=columns)  # type: ignore[arg-type]
+    return DataFrame._from_dict(data=data, columns=columns)
 
 
 def from_dicts(
@@ -170,7 +170,7 @@ def from_records(
 
 
 def from_numpy(
-    data: np.ndarray,
+    data: np.ndarray[Any, Any],
     columns: Sequence[str] | None = None,
     orient: Literal["col", "row"] | None = None,
 ) -> DataFrame:

py-polars/polars/datatypes_constructor.py

@@ -39,8 +39,11 @@
 except ImportError:
     _NUMPY_AVAILABLE = False
 
+
 if not _DOCUMENTING:
-    _POLARS_TYPE_TO_CONSTRUCTOR: dict[PolarsDataType, Callable] = {
+    _POLARS_TYPE_TO_CONSTRUCTOR: dict[
+        PolarsDataType, Callable[[str, Sequence[Any], bool], PySeries]
+    ] = {
         Float32: PySeries.new_opt_f32,
         Float64: PySeries.new_opt_f64,
         Int8: PySeries.new_opt_i8,
@@ -91,7 +94,7 @@ def polars_type_to_constructor(
     }
 
 
-def numpy_type_to_constructor(dtype: type[np.dtype]) -> Callable[..., PySeries]:
+def numpy_type_to_constructor(dtype: type[np.dtype[Any]]) -> Callable[..., PySeries]:
     """Get the right PySeries constructor for the given Polars dtype."""
     try:
         return _NUMPY_TYPE_TO_CONSTRUCTOR[dtype]

py-polars/polars/internals/anonymous_scan.py

@@ -93,7 +93,7 @@ def _scan_ipc_impl(uri: str, with_columns: list[str] | None) -> pli.DataFrame:
 
 def _scan_ipc_fsspec(
     file: str,
-    storage_options: dict | None = None,
+    storage_options: dict[str, object] | None = None,
 ) -> pli.LazyFrame:
     func = partial(_scan_ipc_impl, file)
     func_serialized = pickle.dumps(func)
@@ -124,7 +124,7 @@ def _scan_parquet_impl(uri: str, with_columns: list[str] | None) -> pli.DataFram
 
 def _scan_parquet_fsspec(
     file: str,
-    storage_options: dict | None = None,
+    storage_options: dict[str, object] | None = None,
 ) -> pli.LazyFrame:
     func = partial(_scan_parquet_impl, file)
     func_serialized = pickle.dumps(func)

py-polars/polars/internals/construction.py

@@ -4,7 +4,7 @@
 from contextlib import suppress
 from datetime import date, datetime, time, timedelta
 from itertools import zip_longest
-from typing import TYPE_CHECKING, Any, Iterable, Sequence
+from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence
 
 from polars import internals as pli
 from polars.datatypes import (
@@ -91,7 +91,10 @@ def arrow_to_pyseries(name: str, values: pa.Array, rechunk: bool = True) -> PySe
 
 
 def numpy_to_pyseries(
-    name: str, values: np.ndarray, strict: bool = True, nan_to_null: bool = False
+    name: str,
+    values: np.ndarray[Any, Any],
+    strict: bool = True,
+    nan_to_null: bool = False,
 ) -> PySeries:
     """Construct a PySeries from a numpy array."""
     if not values.flags["C_CONTIGUOUS"]:
@@ -411,7 +414,8 @@ def _unpack_columns(
 
 
 def dict_to_pydf(
-    data: dict[str, Sequence[Any]], columns: ColumnsType | None = None
+    data: Mapping[str, Sequence[object] | Mapping[str, Sequence[object]]],
+    columns: ColumnsType | None = None,
 ) -> PyDataFrame:
     """Construct a PyDataFrame from a dictionary of sequences."""
     if columns is not None:
@@ -520,7 +524,7 @@ def sequence_to_pydf(
 
 
 def numpy_to_pydf(
-    data: np.ndarray,
+    data: np.ndarray[Any, Any],
     columns: ColumnsType | None = None,
     orient: Literal["col", "row"] | None = None,
 ) -> PyDataFrame:

py-polars/polars/internals/expr.py

@@ -1757,7 +1757,7 @@ def sort_by(
 
         return wrap_expr(self._pyexpr.sort_by(by, reverse))
 
-    def take(self, index: List[int] | Expr | pli.Series | np.ndarray) -> Expr:
+    def take(self, index: List[int] | Expr | pli.Series | np.ndarray[Any, Any]) -> Expr:
         """
         Take values by index.
 
@@ -1842,7 +1842,7 @@ def shift(self, periods: int = 1) -> Expr:
         return wrap_expr(self._pyexpr.shift(periods))
 
     def shift_and_fill(
-        self, periods: int, fill_value: int | float | bool | str | Expr | list
+        self, periods: int, fill_value: int | float | bool | str | Expr | list[Any]
     ) -> Expr:
         """
         Shift the values by a given period and fill the parts that will be empty due to

py-polars/polars/internals/frame.py

@@ -266,7 +266,7 @@ def __init__(
         data: (
             dict[str, Sequence[Any]]
             | Sequence[Any]
-            | np.ndarray
+            | np.ndarray[Any, Any]
             | pa.Table
             | pd.DataFrame
             | pli.Series
@@ -341,7 +341,7 @@ def _from_dicts(
     @classmethod
     def _from_dict(
         cls: type[DF],
-        data: dict[str, Sequence[Any]],
+        data: Mapping[str, Sequence[object] | Mapping[str, Sequence[object]]],
         columns: Sequence[str] | None = None,
     ) -> DF:
         """
@@ -395,7 +395,7 @@ def _from_records(
     @classmethod
     def _from_numpy(
         cls: type[DF],
-        data: np.ndarray,
+        data: np.ndarray[Any, Any],
         columns: Sequence[str] | None = None,
         orient: Literal["col", "row"] | None = None,
     ) -> DF:
@@ -1496,7 +1496,7 @@ def to_parquet(
             **kwargs,
         )
 
-    def to_numpy(self) -> np.ndarray:
+    def to_numpy(self) -> np.ndarray[Any, Any]:
         """
         Convert DataFrame to a 2d numpy array.
         This operation clones data.
@@ -1704,7 +1704,9 @@ def _pos_idx(self, idx: int, dim: int) -> int:
         else:
             return self.shape[dim] + idx
 
-    def _pos_idxs(self, idxs: np.ndarray | pli.Series, dim: int) -> pli.Series:
+    def _pos_idxs(
+        self, idxs: np.ndarray[Any, Any] | pli.Series, dim: int
+    ) -> pli.Series:
         # pl.UInt32 (polars) or pl.UInt64 (polars_u64_idx).
         idx_type = get_idx_type()
 
@@ -1784,7 +1786,7 @@ def __getitem__(self: DF, item: str) -> pli.Series:
     def __getitem__(
         self: DF,
         item: int
-        | np.ndarray
+        | np.ndarray[Any, Any]
         | pli.Expr
         | list[pli.Expr]
         | MultiColSelector
@@ -1814,7 +1816,7 @@ def __getitem__(
         item: (
             str
             | int
-            | np.ndarray
+            | np.ndarray[Any, Any]
             | pli.Expr
             | list[pli.Expr]
             | MultiColSelector
@@ -1981,7 +1983,7 @@ def __getitem__(
         )
 
     def __setitem__(
-        self, key: str | list | tuple[Any, str | int], value: Any
+        self, key: str | list[Any] | tuple[Any, str | int], value: Any
     ) -> None:  # pragma: no cover
         warnings.warn(
             "setting a DataFrame by indexing is deprecated; Consider using"
@@ -3218,7 +3220,7 @@ def groupby_rolling(
         return RollingGroupBy(self, index_column, period, offset, closed, by)
 
     def groupby_dynamic(
-        self,
+        self: DF,
         index_column: str,
         every: str,
         period: str | None = None,
@@ -3227,7 +3229,7 @@ def groupby_dynamic(
         include_boundaries: bool = False,
         closed: ClosedWindow = "right",
         by: str | list[str] | pli.Expr | list[pli.Expr] | None = None,
-    ) -> DynamicGroupBy:
+    ) -> DynamicGroupBy[DF]:
         """
         Group based on a time value (or index value of type Int32, Int64).
 
@@ -3633,7 +3635,7 @@ def upsample(
 
     @deprecated_alias(df="other")
     def join_asof(
-        self: DF,
+        self,
         other: DataFrame,
         left_on: str | None = None,
         right_on: str | None = None,
@@ -3646,7 +3648,7 @@ def join_asof(
         tolerance: str | int | float | None = None,
         allow_parallel: bool = True,
         force_parallel: bool = False,
-    ) -> DF:
+    ) -> DataFrame:
         """
         Perform an asof join. This is similar to a left-join except that we
         match on nearest key rather than equal keys.
@@ -3777,7 +3779,7 @@ def join_asof(
 
     @deprecated_alias(df="other")
     def join(
-        self: DF,
+        self,
         other: DataFrame,
         left_on: str | pli.Expr | list[str | pli.Expr] | None = None,
         right_on: str | pli.Expr | list[str | pli.Expr] | None = None,
@@ -3787,7 +3789,7 @@ def join(
         asof_by: str | list[str] | None = None,
         asof_by_left: str | list[str] | None = None,
         asof_by_right: str | list[str] | None = None,
-    ) -> DF:
+    ) -> DataFrame:
         """
         Join in SQL-like fashion.
 
@@ -5978,7 +5980,7 @@ def row(self, index: int) -> tuple[Any]:
         """
         return self._df.row_tuple(index)
 
-    def rows(self) -> list[tuple]:
+    def rows(self) -> list[tuple[object, ...]]:
         """
         Convert columnar data to rows as python tuples.
 

py-polars/polars/internals/lazy_frame.py

@@ -170,7 +170,7 @@ def scan_parquet(
         rechunk: bool = True,
         row_count_name: str | None = None,
         row_count_offset: int = 0,
-        storage_options: dict | None = None,
+        storage_options: dict[str, object] | None = None,
         low_memory: bool = False,
     ) -> LDF:
         """
@@ -213,7 +213,7 @@ def scan_ipc(
         rechunk: bool = True,
         row_count_name: str | None = None,
         row_count_offset: int = 0,
-        storage_options: dict | None = None,
+        storage_options: dict[str, object] | None = None,
         memory_map: bool = True,
     ) -> LDF:
         """

py-polars/polars/internals/lazy_functions.py

@@ -607,10 +607,7 @@ def tail(column: str | pli.Series, n: int | None = None) -> pli.Expr | pli.Serie
     return col(column).tail(n)
 
 
-def lit(
-    value: float | int | str | date | datetime | pli.Series | np.ndarray | Any | None,
-    dtype: type[DataType] | None = None,
-) -> pli.Expr:
+def lit(value: Any, dtype: type[DataType] | None = None) -> pli.Expr:
     """
     Return an expression representing a literal value.
 
@@ -682,11 +679,13 @@ def lit(
 
     if dtype:
         return pli.wrap_expr(pylit(value)).cast(dtype)
-    # numpy literals like np.float32(0)
-    # have an item
-    if hasattr(value, "item"):
-        value = value.item()  # type: ignore[union-attr]
-    return pli.wrap_expr(pylit(value))
+
+    try:
+        # numpy literals like np.float32(0) have an item
+        item = value.item()
+    except AttributeError:
+        item = value
+    return pli.wrap_expr(pylit(item))
 
 
 def spearman_rank_corr(