Add a generic Series type for polars

pandera/api/polars/model.py

@@ -21,6 +21,7 @@
 from pandera.engines import polars_engine as pe
 from pandera.errors import SchemaInitError
 from pandera.typing import AnnotationInfo
+from pandera.typing.polars import Series
 
 
 class DataFrameModel(_DataFrameModel[pl.LazyFrame, DataFrameSchema]):
@@ -74,6 +75,7 @@
             if (
                 annotation.origin is None
                 or isinstance(annotation.origin, pl.datatypes.DataTypeClass)
+                or annotation.origin is Series
                 or engine_dtype
             ):
                 if check_name is False:
@@ -94,19 +96,9 @@
                 columns[field_name] = Column(**column_kwargs)
 
             else:
-                origin_name = (
-                    f"{annotation.origin.__module__}."
-                    f"{annotation.origin.__name__}"
-                )
-                msg = (
-                    " Series[TYPE] annotations are not supported for polars. "
-                    "Use the bare TYPE directly"
-                    if origin_name == "pandera.typing.pandas.Series"
-                    else ""
-                )
                 raise SchemaInitError(
                     f"Invalid annotation '{field_name}: "
-                    f"{annotation.raw_annotation}'.{msg}"
+                    f"{annotation.raw_annotation}'."
                 )
 
         return columns

pandera/typing/polars.py

@@ -4,7 +4,11 @@
 
 from packaging import version
 
-from pandera.typing.common import DataFrameBase, DataFrameModel
+from pandera.typing.common import (
+    DataFrameBase,
+    DataFrameModel,
+    SeriesBase,
+)
 
 
 try:
@@ -35,3 +39,11 @@
 
         *new in 0.19.0*
         """
+
+    # pylint: disable=too-few-public-methods
+    class Series(SeriesBase, pl.Series, Generic[T]):
+        """
+        Pandera generic for pl.Series, only used for type annotation.
+
+        *new in 0.19.0*
+        """

tests/polars/test_polars_dataframe_generic.py

@@ -4,7 +4,24 @@
 import pytest
 
 import pandera.polars as pa
-from pandera.typing.polars import LazyFrame
+from pandera.typing.polars import LazyFrame, Series
+
+
+def test_series_annotation():
+    class Model(pa.DataFrameModel):
+        col1: Series[pl.Int64]
+
+    data = pl.LazyFrame(
+        {
+            "col1": [1, 2, 3],
+        }
+    )
+
+    assert data.collect().equals(Model.validate(data).collect())
+
+    invalid_data = data.cast({"col1": pl.Float64})
+    with pytest.raises(pa.errors.SchemaError):
+        Model.validate(invalid_data).collect()
 
 
 def test_lazyframe_generic_simple():