Repository: unionai-oss/pandera

Issue:
pandas_dtype should accept built-in python types and numpy types
fixes #243. currently, only pandas dtypes, numpy type string aliases, and the `pandera` types can be accepted in the SeriesSchemaBase `pandas_dtype` argument. For ease of use, allow the use of python built-in types and numpy dtypes (not just string aliases) to specify the type of columns


Discussion:

Resolving patch:
diff --git a/pandera/dtypes.py b/pandera/dtypes.py
--- a/pandera/dtypes.py
+++ b/pandera/dtypes.py
@@ -1,7 +1,6 @@
 """Schema datatypes."""
 
 from enum import Enum
-from typing import Union
 
 import numpy as np
 import pandas as pd
@@ -63,8 +62,16 @@ class PandasDtype(Enum):
     2    c
     dtype: object
 
-    You can also directly use the string alias for each data-type in the
-    schema definition:
+    Alternatively, you can use built-in python scalar types for integers,
+    floats, booleans, and strings:
+
+    >>> pa.SeriesSchema(int).validate(pd.Series([1, 2, 3]))
+    0    1
+    1    2
+    2    3
+    dtype: int64
+
+    You can also use the pandas string aliases in the schema definition:
 
     >>> pa.SeriesSchema("int").validate(pd.Series([1, 2, 3]))
     0    1
@@ -168,7 +175,7 @@ def from_str_alias(cls, str_alias: str) -> "PandasDtype":
         }.get(str_alias)
 
         if pandas_dtype is None:
-            raise ValueError(
+            raise TypeError(
                 "pandas dtype string alias '%s' not recognized" %
                 str_alias
             )
@@ -176,8 +183,7 @@ def from_str_alias(cls, str_alias: str) -> "PandasDtype":
         return pandas_dtype
 
     @classmethod
-    def from_pandas_api_type(
-            cls, pandas_api_type: str) -> Union["PandasDtype", None]:
+    def from_pandas_api_type(cls, pandas_api_type: str) -> "PandasDtype":
         """Get PandasDtype enum from pandas api type.
 
         :param pandas_api_type: string output from
@@ -187,7 +193,7 @@ def from_pandas_api_type(
         if pandas_api_type.startswith("mixed"):
             return cls.Object
 
-        return {
+        pandas_dtype = {
             "string": cls.String,
             "floating": cls.Float,
             "integer": cls.Int,
@@ -199,6 +205,36 @@ def from_pandas_api_type(
             "timedelta": cls.Timedelta,
         }.get(pandas_api_type)
 
+        if pandas_dtype is None:
+            raise TypeError(
+                "pandas api type '%s' not recognized" %
+                pandas_api_type
+            )
+
+        return pandas_dtype
+
+    @classmethod
+    def from_python_type(cls, python_type: type) -> "PandasDtype":
+        """Get PandasDtype enum from built-in python type.
+
+        :param python_type: built-in python type. Allowable types are:
+            str, int, float, and bool.
+        """
+        pandas_dtype = {
+            bool: cls.Bool,
+            str: cls.String,
+            int: cls.Int,
+            float: cls.Float,
+        }.get(python_type)
+
+        if pandas_dtype is None:
+            raise TypeError(
+                "python type '%s' not recognized as pandas data type" %
+                python_type
+            )
+
+        return pandas_dtype
+
     def __eq__(self, other):
         # pylint: disable=comparison-with-callable,too-many-return-statements
         # see https://github.com/PyCQA/pylint/issues/2306
diff --git a/pandera/schema_components.py b/pandera/schema_components.py
--- a/pandera/schema_components.py
+++ b/pandera/schema_components.py
@@ -7,9 +7,11 @@
 import numpy as np
 import pandas as pd
 
-from . import errors, dtypes
+from . import errors
 from .dtypes import PandasDtype
-from .schemas import DataFrameSchema, SeriesSchemaBase, CheckList
+from .schemas import (
+    DataFrameSchema, SeriesSchemaBase, CheckList, PandasDtypeInputTypes
+)
 
 
 def _is_valid_multiindex_tuple_str(x: Tuple[Any]) -> bool:
@@ -22,8 +24,7 @@ class Column(SeriesSchemaBase):
 
     def __init__(
             self,
-            pandas_dtype: Union[
-                str, PandasDtype, dtypes.PandasExtensionType] = None,
+            pandas_dtype: PandasDtypeInputTypes = None,
             checks: CheckList = None,
             nullable: bool = False,
             allow_duplicates: bool = True,
@@ -241,13 +242,13 @@ class Index(SeriesSchemaBase):
 
     def __init__(
             self,
-            pandas_dtype: Union[
-                str, PandasDtype, dtypes.PandasExtensionType] = None,
+            pandas_dtype: PandasDtypeInputTypes = None,
             checks: CheckList = None,
             nullable: bool = False,
             allow_duplicates: bool = True,
             coerce: bool = False,
             name: str = None) -> None:
+        # pylint: disable=useless-super-delegation
         """Create Index validator.
 
         :param pandas_dtype: datatype of the column. A ``PandasDtype`` for
diff --git a/pandera/schemas.py b/pandera/schemas.py
--- a/pandera/schemas.py
+++ b/pandera/schemas.py
@@ -11,7 +11,7 @@
 
 from . import errors, constants, dtypes
 from .checks import Check
-from .dtypes import PandasDtype
+from .dtypes import PandasDtype, PandasExtensionType
 from .error_formatters import (
     format_generic_error_message, format_vectorized_error_message,
     reshape_failure_cases, scalar_failure_case,
@@ -29,6 +29,8 @@
     ]
 ]
 
+PandasDtypeInputTypes = Union[str, type, PandasDtype, PandasExtensionType]
+
 
 def _inferred_schema_guard(method):
     """
@@ -614,9 +616,7 @@ class SeriesSchemaBase():
 
     def __init__(
             self,
-            pandas_dtype: Union[
-                str, dtypes.PandasDtype, dtypes.PandasExtensionType
-            ] = None,
+            pandas_dtype: PandasDtypeInputTypes = None,
             checks: CheckList = None,
             nullable: bool = False,
             allow_duplicates: bool = True,
@@ -656,6 +656,12 @@ def __init__(
                 raise errors.SchemaInitError(
                     "Cannot use groupby checks with type %s" % type(self))
 
+        # make sure pandas dtype is valid
+        try:
+            self.dtype
+        except TypeError:
+            raise
+
         # this attribute is not meant to be accessed by users and is explicitly
         # set to True in the case that a schema is created by infer_schema.
         self._IS_INFERRED = False
@@ -750,6 +756,8 @@ def dtype(self) -> Union[str, None]:
         elif isinstance(self._pandas_dtype, str):
             dtype = PandasDtype.from_str_alias(  # type: ignore
                 self._pandas_dtype).str_alias
+        elif isinstance(self._pandas_dtype, type):
+            dtype = PandasDtype.from_python_type(self._pandas_dtype).str_alias
         elif isinstance(self._pandas_dtype, dtypes.PandasDtype):
             dtype = self._pandas_dtype.str_alias
         else:
@@ -845,8 +853,6 @@ def validate(
 
         series_dtype = series.dtype
         if self._nullable:
-            # currently, to handle null cases drop null values before passing
-            # into checks.
             series_no_nans = series.dropna()
             if self.dtype in dtypes.NUMPY_NONNULLABLE_INT_DTYPES:
                 _series = series_no_nans.astype(self.dtype)


Test patch:
diff --git a/tests/test_dtypes.py b/tests/test_dtypes.py
--- a/tests/test_dtypes.py
+++ b/tests/test_dtypes.py
@@ -163,7 +163,7 @@ def test_pandas_nullable_int_dtype(dtype, coerce):
 @pytest.mark.parametrize("str_alias", ["foo", "bar", "baz", "asdf", "qwerty"])
 def test_unrecognized_str_aliases(str_alias):
     """Test that unrecognized string aliases are supported."""
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         PandasDtype.from_str_alias(str_alias)
 
 
@@ -294,3 +294,63 @@ def test_pandas_extension_types():
         series_kwargs = {} if series_kwargs is None else series_kwargs
         series_schema = SeriesSchema(pandas_dtype=dtype, **series_kwargs)
         assert isinstance(series_schema.validate(data), pd.Series)
+
+
+def test_python_builtin_types():
+    """Test support python data types can be used for validation."""
+    schema = DataFrameSchema({
+        "int_col": Column(int),
+        "float_col": Column(float),
+        "str_col": Column(str),
+        "bool_col": Column(bool),
+    })
+    df = pd.DataFrame({
+        "int_col": [1, 2, 3],
+        "float_col": [1., 2., 3.],
+        "str_col": list("abc"),
+        "bool_col": [True, False, True],
+    })
+    assert isinstance(schema(df), pd.DataFrame)
+    assert schema.dtype["int_col"] == PandasDtype.Int.str_alias
+    assert schema.dtype["float_col"] == PandasDtype.Float.str_alias
+    assert schema.dtype["str_col"] == PandasDtype.String.str_alias
+    assert schema.dtype["bool_col"] == PandasDtype.Bool.str_alias
+
+
+@pytest.mark.parametrize("python_type", [list, dict, set])
+def test_python_builtin_types_not_supported(python_type):
+    """Test unsupported python data types raise a type error."""
+    with pytest.raises(TypeError):
+        Column(python_type)
+
+
+@pytest.mark.parametrize(
+    "pandas_api_type,pandas_dtype", [
+        ["string", PandasDtype.String],
+        ["floating", PandasDtype.Float],
+        ["integer", PandasDtype.Int],
+        ["categorical", PandasDtype.Category],
+        ["boolean", PandasDtype.Bool],
+        ["datetime64", PandasDtype.DateTime],
+        ["datetime", PandasDtype.DateTime],
+        ["timedelta64", PandasDtype.Timedelta],
+        ["timedelta", PandasDtype.Timedelta],
+    ]
+)
+def test_pandas_api_types(pandas_api_type, pandas_dtype):
+    """Test pandas api type conversion."""
+    assert PandasDtype.from_pandas_api_type(pandas_api_type) is pandas_dtype
+
+
+@pytest.mark.parametrize(
+    "invalid_pandas_api_type", [
+        "foo",
+        "bar",
+        "baz",
+        "this is not a type",
+    ]
+)
+def test_pandas_api_type_exception(invalid_pandas_api_type):
+    """Test unsupported values for pandas api type conversion."""
+    with pytest.raises(TypeError):
+        PandasDtype.from_pandas_api_type(invalid_pandas_api_type)
diff --git a/tests/test_schema_components.py b/tests/test_schema_components.py
--- a/tests/test_schema_components.py
+++ b/tests/test_schema_components.py
@@ -474,7 +474,7 @@ def test_column_type_can_be_set():
     assert column_a.dtype == changed_type.str_alias
 
     for invalid_dtype in ("foobar", "bar"):
-        with pytest.raises(ValueError):
+        with pytest.raises(TypeError):
             column_a.pandas_dtype = invalid_dtype
 
     for invalid_dtype in (1, 2.2, ["foo", 1, 1.1], {"b": 1}):