refine things

Signed-off-by: Matt Richards <[email protected]>

Author: m-richards

pandera/backends/polars/checks.py

@@ -45,7 +45,9 @@ def preprocess(self, check_obj: pl.LazyFrame, key: Optional[str]):
         # for the index to groupby on. Right now grouping by the index is not allowed.
         return check_obj
 
-    def apply(self, check_obj: PolarsData | AllColumnsPolarsCheckData):
+    def apply(
+        self, check_obj: PolarsData | AllColumnsPolarsCheckData
+    ) -> bool | pl.LazyFrame:
         """Apply the check function to a check object."""
         if self.check.element_wise:
             selector = pl.col(check_obj.key or "*")
@@ -75,11 +77,13 @@ def apply(self, check_obj: PolarsData | AllColumnsPolarsCheckData):
 
         return out
 
-    def postprocess(self, check_obj, check_output):
+    def postprocess(
+        self, check_obj: PolarsData | AllColumnsPolarsCheckData, check_output
+    ):
         """Postprocesses the result of applying the check function."""
-        if isinstance(check_obj, PolarsData) and isinstance(
-            check_output, pl.LazyFrame
-        ):
+        if isinstance(
+            check_obj, PolarsData | AllColumnsPolarsCheckData
+        ) and isinstance(check_output, pl.LazyFrame):
             return self.postprocess_lazyframe_output(check_obj, check_output)
         elif isinstance(check_output, bool):
             return self.postprocess_bool_output(check_obj, check_output)
@@ -89,7 +93,7 @@ def postprocess(self, check_obj, check_output):
 
     def postprocess_lazyframe_output(
         self,
-        check_obj: PolarsData,
+        check_obj: PolarsData | AllColumnsPolarsCheckData,
         check_output: pl.LazyFrame,
     ) -> CheckResult:
         """Postprocesses the result of applying the check function."""
@@ -114,7 +118,7 @@ def postprocess_lazyframe_output(
 
     def postprocess_bool_output(
         self,
-        check_obj: PolarsData,
+        check_obj: PolarsData | AllColumnsPolarsCheckData,
         check_output: bool,
     ) -> CheckResult:
         """Postprocesses the result of applying the check function."""

pandera/engines/polars_engine.py

@@ -16,6 +16,7 @@
     Type,
     Union,
     TypedDict,
+    overload,
 )
 
 import polars as pl
@@ -578,6 +579,24 @@ class Array(DataType):
 
     type = pl.Array
 
+    @overload
+    def __init__(
+        self,
+        inner: Literal[None] = ...,
+        shape: Literal[None] = ...,
+        *,
+        width: Literal[None] = ...,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        inner: PolarsDataType = ...,
+        shape: Union[int, Tuple[int, ...], None] = ...,
+        *,
+        width: Optional[int] = ...,
+    ) -> None: ...
+
     def __init__(  # pylint:disable=super-init-not-called
         self,
         inner: Optional[PolarsDataType] = None,
@@ -593,8 +612,7 @@ def __init__(  # pylint:disable=super-init-not-called
             kwargs["shape"] = width
         elif shape is not None:
             kwargs["shape"] = shape
-
-        if inner or shape or width:
+        if inner:
             object.__setattr__(self, "type", pl.Array(inner=inner, **kwargs))
 
     @classmethod

tests/polars/test_polars_dtypes.py

@@ -13,7 +13,7 @@
 from polars.testing.parametric import dataframes
 
 import pandera.errors
-from pandera.api.polars.types import PolarsData
+from pandera.api.polars.types import AllColumnsPolarsCheckData
 from pandera.api.polars.utils import get_lazyframe_column_dtypes
 from pandera.constants import CHECK_OUTPUT_KEY
 from pandera.engines import polars_engine as pe
@@ -98,7 +98,9 @@ def test_coerce_no_cast(dtype, data):
         pl.enable_string_cache()
     pandera_dtype = dtype()
     df = data.draw(get_dataframe_strategy(type_=pandera_dtype.type))
-    coerced = pandera_dtype.coerce(data_container=PolarsData(df))
+    coerced = pandera_dtype.coerce(
+        data_container=AllColumnsPolarsCheckData(df)
+    )
     assert_frame_equal(df, coerced)
 
 
@@ -317,7 +319,9 @@ def test_polars_object_coercible(to_dtype, container, result):
     Test that polars_object_coercible can detect that a polars object is
     coercible or not.
     """
-    is_coercible = polars_object_coercible(PolarsData(container), to_dtype)
+    is_coercible = polars_object_coercible(
+        AllColumnsPolarsCheckData(container), to_dtype
+    )
     assert_frame_equal(is_coercible, result)
 
 
@@ -439,12 +443,14 @@ def test_polars_nested_dtypes_try_coercion(
     data,
 ):
     pandera_dtype = pe.Engine.dtype(coercible_dtype)
-    coerced_data = pandera_dtype.try_coerce(PolarsData(data))
+    coerced_data = pandera_dtype.try_coerce(AllColumnsPolarsCheckData(data))
     assert coerced_data.collect().equals(data.collect())
 
     # coercing data with invalid type should raise an error
     try:
-        pe.Engine.dtype(noncoercible_dtype).try_coerce(PolarsData(data))
+        pe.Engine.dtype(noncoercible_dtype).try_coerce(
+            AllColumnsPolarsCheckData(data)
+        )
     except pandera.errors.ParserError as exc:
         col = pl.col(exc.failure_cases.columns[0])
         assert exc.failure_cases.select(col).equals(data.collect())