Repository: unionai-oss/pandera (Python)

Issue:
Parametrized type annotations are broken for polars DataFrameModels
# Description

Pandera DataFrameModels do not support parameterized types for polars, while DataFrameSchemas do.

- [x] I have checked that this issue has not already been reported.
- [x] I have confirmed this bug exists on the latest version of pandera.
- [ ] (optional) I have confirmed this bug exists on the master branch of pandera.

## Example

Here is an example of a working DataFrameSchema and several variations of broken DataFrameModels.

```python
from typing import Annotated

import pandera.polars as pa
import polars as pl

from pandera.typing import Series
from pandera.errors import SchemaInitError


df = pl.DataFrame({
    "id": [1, 2, 3],
    "lists": [["a"], ["a", "b"], ["a", "b", "c"]],
})


# works!
schema = pa.DataFrameSchema(
    columns={
        "id": pa.Column(int),
        "lists": pa.Column(list[str]),
    }
)
schema.validate(df)
print("DataFrameSchema validation passed")


class Lists(pa.DataFrameModel):
    """Most basic, expected form given the working schema above."""
    id: int
    lists: list[str]


try:
    Lists.validate(df)
except SchemaInitError as e:
    print("\nLists validation failed")
    print(e)
else:
    print("\nLists validation passed")


class ListsSeries(pa.DataFrameModel):
    """Using series as a wrapper around basic data types like the id column here
    will not work. Examples of this appear in the pandera documentation.
    https://pandera.readthedocs.io/en/latest/dataframe_models.html#dtype-aliases
    """
    id: Series[int]
    lists: Series[list[str]]


try:
    ListsSeries.validate(df)
except SchemaInitError as e:
    print("\nListsSeries validation failed")
    print(e)
else:
    print("\nListsSeries validation passed")


class AlternateListsSeries(pa.DataFrameModel):
    """Demonstrating using Series as a type wrapper around only lists to avoid
    the initialization error on id."""
    id: int
    lists: Series[list[str]]


try:
    AlternateListsSeries.validate(df)
except SchemaInitError as e:
    print("\nAlternateListsSeries validation failed")
    print(e)
else:
    print("\nAlternateListsSeries validation passed")


class ListsAnnotated(pa.DataFrameModel):
    """Parameterized form using Annotated as suggested at
    https://pandera.readthedocs.io/en/latest/polars.html#nested-types
    """
    id: int
    lists: Series[Annotated[list, str]]


try:
    ListsAnnotated.validate(df)
except TypeError as e:
    print("\nListsAnnotated validation failed")
    print(e)
else:
    print("\nListsAnnotated validation passed")


class ListsAnnotatedStr(pa.DataFrameModel):
    """Alternate parameterized form using Annotated as seen in the examples here:
    https://pandera.readthedocs.io/en/latest/dataframe_models.html#annotated
    """
    id: int
    lists: Series[Annotated[list, "str"]]

try:
    ListsAnnotatedStr.validate(df)
except TypeError as e:
    print("\nListsAnnotatedStr validation failed")
    print(e)
else:
    print("\nListsAnnotatedStr validation passed")
```

When run with the following python / library versions:
* python==3.11.9
* polars==0.20.19
* pandera[polars]==0.19.0b1

the above script produces:
```
DataFrameSchema validation passed

Lists validation failed
Invalid annotation 'lists: list[str]'

ListsSeries validation failed
Invalid annotation 'id: pandera.typing.pandas.Series[int]'

AlternateListsSeries validation failed
Invalid annotation 'lists: pandera.typing.pandas.Series[list[str]]'

ListsAnnotated validation failed
Annotation 'Annotated' requires all positional arguments ['args', 'kwargs'].

ListsAnnotatedStr validation failed
Annotation 'Annotated' requires all positional arguments ['args', 'kwargs'].
```

## Expected behavior

I would expect any column types that are valid to pass to DataFrameSchema's constructor to also be valid as annotations for DataFrameModel.

## Desktop (please complete the following information):

 - OS: MacOS 14.4.1 (M2 Max)
 - Browser: Chrome
 - Version: 0.19.0b1



Fix (patch):
diff --git a/pandera/api/polars/model.py b/pandera/api/polars/model.py
index 48eeb482b..c51959ee1 100644
--- a/pandera/api/polars/model.py
+++ b/pandera/api/polars/model.py
@@ -1,7 +1,6 @@
 """Class-based api for polars models."""
 
 from typing import (
-    Any,
     Dict,
     List,
     Tuple,
@@ -19,6 +18,7 @@
 from pandera.api.polars.container import DataFrameSchema
 from pandera.api.polars.components import Column
 from pandera.api.polars.model_config import BaseConfig
+from pandera.engines import polars_engine as pe
 from pandera.errors import SchemaInitError
 from pandera.typing import AnnotationInfo
 
@@ -52,24 +52,29 @@ def _build_columns(  # pylint:disable=too-many-locals
             field_name = field.name
             check_name = getattr(field, "check_name", None)
 
-            if annotation.metadata:
-                if field.dtype_kwargs:
-                    raise TypeError(
-                        "Cannot specify redundant 'dtype_kwargs' "
-                        + f"for {annotation.raw_annotation}."
-                        + "\n Usage Tip: Drop 'typing.Annotated'."
-                    )
-                dtype_kwargs = get_dtype_kwargs(annotation)
-                dtype = annotation.arg(**dtype_kwargs)  # type: ignore
-            elif annotation.default_dtype:
-                dtype = annotation.default_dtype
-            else:
-                dtype = annotation.arg
-
-            dtype = None if dtype is Any else dtype
-
-            if annotation.origin is None or isinstance(
-                annotation.origin, pl.datatypes.DataTypeClass
+            engine_dtype = None
+            try:
+                engine_dtype = pe.Engine.dtype(annotation.raw_annotation)
+                dtype = engine_dtype.type
+            except TypeError as exc:
+                if annotation.metadata:
+                    if field.dtype_kwargs:
+                        raise TypeError(
+                            "Cannot specify redundant 'dtype_kwargs' "
+                            + f"for {annotation.raw_annotation}."
+                            + "\n Usage Tip: Drop 'typing.Annotated'."
+                        ) from exc
+                    dtype_kwargs = get_dtype_kwargs(annotation)
+                    dtype = annotation.arg(**dtype_kwargs)  # type: ignore
+                elif annotation.default_dtype:
+                    dtype = annotation.default_dtype
+                else:
+                    dtype = annotation.arg
+
+            if (
+                annotation.origin is None
+                or isinstance(annotation.origin, pl.datatypes.DataTypeClass)
+                or engine_dtype
             ):
                 if check_name is False:
                     raise SchemaInitError(