implement coercion-on-initialization for DataFrame[SchemaModel] types (…

…#772)

* implement coercion-on-initialization

* pylint

* Update tests/core/test_model.py

Co-authored-by: Matt Richards <45483497+m-richards@users.noreply.github.com>

Co-authored-by: Matt Richards <45483497+m-richards@users.noreply.github.com>

.pylintrc

@@ -25,4 +25,5 @@ disable=
     no-else-return,
     inconsistent-return-statements,
     protected-access,
-    too-many-ancestors
+    too-many-ancestors,
+    too-many-lines

pandera/typing/common.py

@@ -159,14 +159,13 @@ def __setattr__(self, name: str, value: Any) -> None:
 
             # prevent the double validation problem by preventing checks for
             # dataframes with a defined pandera.schema
-            pandera = getattr(self, "pandera")
+            pandera_accessor = getattr(self, "pandera")
             if (
-                pandera.schema is None
-                or pandera.schema != schema_model.to_schema()
+                pandera_accessor.schema is None
+                or pandera_accessor.schema != schema_model.to_schema()
             ):
-                # pylint: disable=self-cls-assignment
-                self = schema_model.validate(self)
-                pandera.add_schema(schema_model.to_schema())
+                pandera_accessor.add_schema(schema_model.to_schema())
+                self.__dict__ = schema_model.validate(self).__dict__
 
 
 # pylint:disable=too-few-public-methods

pandera/typing/dask.py

@@ -1,7 +1,6 @@
 """Pandera type annotations for Dask."""
 
-import inspect
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 from .common import DataFrameBase, IndexBase, SeriesBase
 from .pandas import GenericDtype, Schema
@@ -32,26 +31,6 @@ class DataFrame(DataFrameBase, dd.DataFrame, Generic[T]):
         *new in 0.8.0*
         """
 
-        def __setattr__(self, name: str, value: Any) -> None:
-            object.__setattr__(self, name, value)
-            if name == "__orig_class__":
-                class_args = getattr(self.__orig_class__, "__args__", None)
-                if class_args is not None and any(
-                    x.__name__ == "SchemaModel"
-                    for x in inspect.getmro(class_args[0])
-                ):
-                    schema_model = value.__args__[0]
-
-                # prevent the double validation problem by preventing checks
-                # for dataframes with a defined pandera.schema
-                if (
-                    self.pandera.schema is None
-                    or self.pandera.schema != schema_model.to_schema()
-                ):
-                    # pylint: disable=self-cls-assignment
-                    self.__dict__ = schema_model.validate(self).__dict__
-                    self.pandera.add_schema(schema_model.to_schema())
-
     # pylint:disable=too-few-public-methods
     class Series(SeriesBase, dd.Series, Generic[GenericDtype]):  # type: ignore
         """Representation of pandas.Series, only used for type annotation.

tests/core/test_model.py

@@ -941,6 +941,8 @@ class Child(Base):
 
 
 def test_column_access_regex() -> None:
+    """Test that column regex alias is reflected in schema attribute."""
+
     class Schema(pa.SchemaModel):
         col_regex: Series[str] = pa.Field(alias="column_([0-9])+", regex=True)
 
@@ -965,3 +967,38 @@ class Bar(pa.SchemaModel):
 
     assert Foo.Config.name == "foo"
     assert Bar.Config.name == "Bar"
+
+
+def test_validate_coerce_on_init():
+    """Test that DataFrame[Schema] validates and coerces on initialization."""
+
+    class Schema(pa.SchemaModel):
+        state: Series[str]
+        city: Series[str]
+        price: Series[float] = pa.Field(
+            in_range={"min_value": 5, "max_value": 20}
+        )
+
+        class Config:
+            coerce = True
+
+    class SchemaNoCoerce(Schema):
+        class Config:
+            coerce = False
+
+    raw_data = {
+        "state": ["NY", "FL", "GA", "CA"],
+        "city": ["New York", "Miami", "Atlanta", "San Francisco"],
+        "price": [8, 12, 10, 16],
+    }
+    pandera_validated_df = DataFrame[Schema](raw_data)
+    pandas_df = pd.DataFrame(raw_data)
+    assert pandera_validated_df.equals(Schema.validate(pandas_df))
+    assert isinstance(pandera_validated_df, DataFrame)
+    assert isinstance(pandas_df, pd.DataFrame)
+
+    with pytest.raises(
+        pa.errors.SchemaError,
+        match="^expected series 'price' to have type float64, got int64$",
+    ):
+        DataFrame[SchemaNoCoerce](raw_data)

tests/dask/test_dask.py

@@ -1,6 +1,5 @@
 """ Tests that basic Pandera functionality works for Dask objects. """
 
-
 import dask.dataframe as dd
 import pandas as pd
 import pytest