Bugfix/763 improve type annotations for DataFrameModel.validate (unionai-oss#1905)

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

* trial type annotations

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

* changes in individual api files

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

* pl.dataframe working in local test

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

* older python union compat

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

* try polars in the mypy env on ci

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

* translate toplevel mypy skip into module specific skips

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

* mypy passes

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

* missing line continuation

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

* python 3.8

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

---------

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

Author: m-richards

pandera/api/pandas/model.py

@@ -2,10 +2,11 @@
 
 import copy
 import sys
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 import pandas as pd
 
+from pandera.api.base.schema import BaseSchema
 from pandera.api.checks import Check
 from pandera.api.dataframe.model import DataFrameModel as _DataFrameModel
 from pandera.api.dataframe.model import get_dtype_kwargs
@@ -22,6 +23,7 @@
     AnnotationInfo,
     DataFrame,
 )
+from pandera.utils import docstring_substitution
 
 # if python version is < 3.11, import Self from typing_extensions
 if sys.version_info < (3, 11):
@@ -171,6 +173,26 @@ def _build_columns_index(  # pylint:disable=too-many-locals,too-many-branches
 
         return columns, _build_schema_index(indices, **multiindex_kwargs)
 
+    @classmethod
+    @docstring_substitution(validate_doc=BaseSchema.validate.__doc__)
+    def validate(
+        cls: Type[Self],
+        check_obj: pd.DataFrame,
+        head: Optional[int] = None,
+        tail: Optional[int] = None,
+        sample: Optional[int] = None,
+        random_state: Optional[int] = None,
+        lazy: bool = False,
+        inplace: bool = False,
+    ) -> DataFrame[Self]:
+        """%(validate_doc)s"""
+        return cast(
+            DataFrame[Self],
+            cls.to_schema().validate(
+                check_obj, head, tail, sample, random_state, lazy, inplace
+            ),
+        )
+
     @classmethod
     def to_json_schema(cls):
         """Serialize schema metadata into json-schema format.

pandera/api/polars/model.py

@@ -1,11 +1,13 @@
 """Class-based api for polars models."""
 
 import inspect
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Tuple, Type, cast, Optional, overload, Union
+from typing_extensions import Self
 
 import pandas as pd
 import polars as pl
 
+from pandera.api.base.schema import BaseSchema
 from pandera.api.checks import Check
 from pandera.api.dataframe.model import DataFrameModel as _DataFrameModel
 from pandera.api.dataframe.model import get_dtype_kwargs
@@ -16,7 +18,8 @@
 from pandera.engines import polars_engine as pe
 from pandera.errors import SchemaInitError
 from pandera.typing import AnnotationInfo
-from pandera.typing.polars import Series
+from pandera.typing.polars import Series, LazyFrame, DataFrame
+from pandera.utils import docstring_substitution
 
 
 class DataFrameModel(_DataFrameModel[pl.LazyFrame, DataFrameSchema]):
@@ -109,6 +112,53 @@ def _build_columns(  # pylint:disable=too-many-locals
 
         return columns
 
+    @classmethod
+    @overload
+    def validate(
+        cls: Type[Self],
+        check_obj: pl.DataFrame,
+        head: Optional[int] = None,
+        tail: Optional[int] = None,
+        sample: Optional[int] = None,
+        random_state: Optional[int] = None,
+        lazy: bool = False,
+        inplace: bool = False,
+    ) -> DataFrame[Self]: ...
+
+    @classmethod
+    @overload
+    def validate(
+        cls: Type[Self],
+        check_obj: pl.LazyFrame,
+        head: Optional[int] = None,
+        tail: Optional[int] = None,
+        sample: Optional[int] = None,
+        random_state: Optional[int] = None,
+        lazy: bool = False,
+        inplace: bool = False,
+    ) -> LazyFrame[Self]: ...
+
+    @classmethod
+    @docstring_substitution(validate_doc=BaseSchema.validate.__doc__)
+    def validate(
+        cls: Type[Self],
+        check_obj: Union[pl.LazyFrame, pl.DataFrame],
+        head: Optional[int] = None,
+        tail: Optional[int] = None,
+        sample: Optional[int] = None,
+        random_state: Optional[int] = None,
+        lazy: bool = False,
+        inplace: bool = False,
+    ) -> Union[LazyFrame[Self], DataFrame[Self]]:
+        """%(validate_doc)s"""
+        result = cls.to_schema().validate(
+            check_obj, head, tail, sample, random_state, lazy, inplace
+        )
+        if isinstance(check_obj, pl.LazyFrame):
+            return cast(LazyFrame[Self], result)
+        else:
+            return cast(DataFrame[Self], result)
+
     @classmethod
     def to_json_schema(cls):
         """Serialize schema metadata into json-schema format.

pandera/api/pyspark/model.py

@@ -41,6 +41,7 @@
 from pandera.errors import SchemaInitError
 from pandera.typing import AnnotationInfo
 from pandera.typing.common import DataFrameBase
+from pandera.typing.pyspark import DataFrame
 
 try:
     from typing_extensions import get_type_hints
@@ -300,10 +301,10 @@ def validate(
         random_state: Optional[int] = None,
         lazy: bool = True,
         inplace: bool = False,
-    ) -> Optional[DataFrameBase[TDataFrameModel]]:
+    ) -> DataFrame[TDataFrameModel]:
         """%(validate_doc)s"""
         return cast(
-            DataFrameBase[TDataFrameModel],
+            DataFrame[TDataFrameModel],
             cls.to_schema().validate(
                 check_obj, head, tail, sample, random_state, lazy, inplace
             ),