more dtype support: int and float types (#101)

unionai-oss · Sep 10, 2019 · b01eef0 · b01eef0
1 parent ffd57d8
commit b01eef0
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 6 deletions.
diff --git a/pandera/dtypes.py b/pandera/dtypes.py
@@ -8,7 +8,18 @@ class PandasDtype(Enum):
     DateTime = "datetime64[ns]"
     Category = "category"
     Float = "float64"
+    Float16 = "float16"
+    Float32 = "float32"
+    Float64 = "float64"
     Int = "int64"
+    Int8 = "int8"
+    Int16 = "int16"
+    Int32 = "int32"
+    Int64 = "int64"
+    UInt8 = "uint8"
+    UInt16 = "uint16"
+    UInt32 = "uint32"
+    UInt64 = "uint64"
     Object = "object"
     String = "object"
     Timedelta = "timedelta64[ns]"
@@ -18,7 +29,18 @@ class PandasDtype(Enum):
 DateTime = PandasDtype.DateTime
 Category = PandasDtype.Category
 Float = PandasDtype.Float
+Float16 = PandasDtype.Float16
+Float32 = PandasDtype.Float32
+Float64 = PandasDtype.Float64
 Int = PandasDtype.Int
+Int8 = PandasDtype.Int8
+Int16 = PandasDtype.Int16
+Int32 = PandasDtype.Int32
+Int64 = PandasDtype.Int64
+UInt8 = PandasDtype.UInt8
+UInt16 = PandasDtype.UInt16
+UInt32 = PandasDtype.UInt32
+UInt64 = PandasDtype.UInt64
 Object = PandasDtype.Object
 String = PandasDtype.String
 Timedelta = PandasDtype.Timedelta
diff --git a/tests/test_pandera.py b/tests/test_pandera.py
@@ -8,6 +8,7 @@
 from pandera import Column, DataFrameSchema, Index, MultiIndex, \
     SeriesSchema, Bool, Category, Check, DateTime, Float, Int, Object, \
     String, Timedelta, check_input, check_output, Hypothesis
+from pandera import dtypes
 from scipy import stats
 
 
@@ -384,12 +385,42 @@ def _assert_expectation(result_df):
         transformer.transform_secord_arg_with_dict_getter(None, dataframe))
 
 
-def test_string_dtypes():
-    # TODO: add tests for all datatypes
-    schema = DataFrameSchema(
-        {"col": Column("float64", nullable=True)})
-    df = pd.DataFrame({"col": [np.nan, 1.0, 2.0]})
-    assert isinstance(schema.validate(df), pd.DataFrame)
+def test_dtypes():
+    for dtype in [
+            dtypes.Float,
+            dtypes.Float16,
+            dtypes.Float32,
+            dtypes.Float64]:
+        schema = DataFrameSchema({"col": Column(dtype, nullable=False)})
+        validated_df = schema.validate(
+            pd.DataFrame(
+                {"col": [-123.1, -7654.321, 1.0, 1.1, 1199.51, 5.1, 4.6]},
+                dtype=dtype.value))
+        assert isinstance(validated_df, pd.DataFrame)
+
+    for dtype in [
+            dtypes.Int,
+            dtypes.Int8,
+            dtypes.Int16,
+            dtypes.Int32,
+            dtypes.Int64]:
+        schema = DataFrameSchema({"col": Column(dtype, nullable=False)})
+        validated_df = schema.validate(
+            pd.DataFrame(
+                {"col": [-712, -4, -321, 0, 1, 777, 5, 123, 9000]},
+                dtype=dtype.value))
+        assert isinstance(validated_df, pd.DataFrame)
+
+    for dtype in [
+            dtypes.UInt8,
+            dtypes.UInt16,
+            dtypes.UInt32,
+            dtypes.UInt64]:
+        schema = DataFrameSchema({"col": Column(dtype, nullable=False)})
+        validated_df = schema.validate(
+            pd.DataFrame(
+                {"col": [1, 777, 5, 123, 9000]}, dtype=dtype.value))
+        assert isinstance(validated_df, pd.DataFrame)
 
 
 def test_nullable_int():