Repository: unionai-oss/pandera

Issue:
add coerce option to DataFrameSchema
this should coerce the type of all columns to the specified dtype before validation.

This should raise a warning when `nullable` Int columns are defined... these should
just default to Float


Discussion:

Resolving patch:
diff --git a/pandera/pandera.py b/pandera/pandera.py
--- a/pandera/pandera.py
+++ b/pandera/pandera.py
@@ -120,7 +120,7 @@ def __call__(self, parent_schema, series, index):
 class DataFrameSchema(object):
     """A light-weight pandas DataFrame validator."""
 
-    def __init__(self, columns, index=None, transformer=None):
+    def __init__(self, columns, index=None, transformer=None, coerce=False):
         """Initialize pandas dataframe schema.
 
         Parameters
@@ -134,19 +134,25 @@ def __init__(self, columns, index=None, transformer=None):
             a callable with signature: pandas.DataFrame -> pandas.DataFrame.
             If specified, calling `validate` will verify properties of the
             columns and return the transformed dataframe object.
-
+        coerce : bool
+            whether or not to coerce all of the columns on validation.
         """
         self.index = index
         self.columns = columns
         self.transformer = transformer
+        self.coerce = coerce
 
     def validate(self, dataframe):
-        for c in self.columns:
+        for c, col in self.columns.items():
             if c not in dataframe:
                 raise SchemaError(
                     "column '%s' not in dataframe\n%s" %
                     (c, dataframe.head()))
 
+            # coercing logic
+            if col.coerce or self.coerce:
+                dataframe[c] = col.coerce_dtype(dataframe[c])
+
         schema_arg = [
             col.set_name(col_name) for col_name, col in self.columns.items()]
         if self.index is not None:
@@ -266,7 +272,7 @@ def __repr__(self):
 class Column(SeriesSchemaBase):
 
     def __init__(
-            self, pandas_dtype, checks=None, nullable=False):
+            self, pandas_dtype, checks=None, nullable=False, coerce=False):
         """Initialize column validator object.
 
         Parameters
@@ -281,14 +287,21 @@ def __init__(
             x is assumed to be a pandas.Series object.
         nullable : bool
             Whether or not column can contain null values.
+        coerce : bool
+            Whether or not to coerce the column to the specified pandas_dtype
+            before validation
         """
         super(Column, self).__init__(pandas_dtype, checks, nullable)
         self._name = None
+        self.coerce = coerce
 
     def set_name(self, name):
         self._name = name
         return self
 
+    def coerce_dtype(self, series):
+        return series.astype(self._pandas_dtype.value)
+
     def __call__(self, df):
         if self._name is None:
             raise RuntimeError(


Test patch:
diff --git a/tests/test_pandera.py b/tests/test_pandera.py
--- a/tests/test_pandera.py
+++ b/tests/test_pandera.py
@@ -7,7 +7,7 @@
 from schema import SchemaError
 
 from pandera import Column, DataFrameSchema, Index, PandasDtype, \
-    SeriesSchema, Check, Int, check_input, check_output
+    SeriesSchema, Check, Int, DateTime, String, check_input, check_output
 
 
 def test_column():
@@ -207,3 +207,38 @@ def test_nullable_int():
         "column1": Column(Int, Check(lambda x: x > 0), nullable=True)
     })
     assert isinstance(null_schema.validate(df), pd.DataFrame)
+
+
+def test_coerce_dtype():
+    df = pd.DataFrame({
+        "column1": [10.0, 20.0, 30.0],
+        "column2": ["2018-01-01", "2018-02-01", "2018-03-01"],
+        "column3": [1, 2, 3],
+        "column4": [1., 1., np.nan],
+    })
+    # specify `coerce` at the Column level
+    schema1 = DataFrameSchema({
+        "column1": Column(Int, Check(lambda x: x > 0), coerce=True),
+        "column2": Column(DateTime, coerce=True),
+        "column3": Column(String, coerce=True),
+    })
+    # specify `coerce` at the DataFrameSchema level
+    schema2 = DataFrameSchema({
+        "column1": Column(Int, Check(lambda x: x > 0)),
+        "column2": Column(DateTime),
+        "column3": Column(String),
+    }, coerce=True)
+
+    for schema in [schema1, schema2]:
+        result = schema.validate(df)
+        assert result.column1.dtype == Int.value
+        assert result.column2.dtype == DateTime.value
+        assert result.column3.dtype == String.value
+
+        # make sure that correct error is raised when null values are present
+        # in a float column that's coerced to an int
+        schema = DataFrameSchema({
+            "column4": Column(Int, coerce=True)
+        })
+        with pytest.raises(ValueError):
+            schema.validate(df)