feat: add overload. (#1823)

* feat: add overload.

Signed-off-by: yassun7010 <[email protected]>

* run pre-commit

Signed-off-by: yassun7010 <[email protected]>

* test: add test cases.

Signed-off-by: yassun7010 <[email protected]>

* fix: add pylint disable.

Signed-off-by: yassun7010 <[email protected]>

* fix: run black format.

Signed-off-by: yassun7010 <[email protected]>

---------

Signed-off-by: yassun7010 <[email protected]>

Author: yassun7010

pandera/dtypes.py

@@ -18,6 +18,8 @@
     Union,
 )
 
+from typing_extensions import overload
+
 try:
     # python 3.8+
     from typing import Literal  # type: ignore[attr-defined]
@@ -89,6 +91,19 @@ def __hash__(self) -> int:
 _DataTypeClass = Type[_Dtype]
 
 
+@overload
+def immutable(
+    pandera_dtype_cls: _DataTypeClass,  # pylint: disable=W0613
+    **dataclass_kwargs: Any,  # pylint: disable=W0613
+) -> _DataTypeClass: ...
+
+
+@overload
+def immutable(
+    **dataclass_kwargs: Any,  # pylint: disable=W0613
+) -> Callable[[_DataTypeClass], _DataTypeClass]: ...
+
+
 def immutable(
     pandera_dtype_cls: Optional[_DataTypeClass] = None, **dataclass_kwargs: Any
 ) -> Union[_DataTypeClass, Callable[[_DataTypeClass], _DataTypeClass]]:

tests/core/test_model.py

@@ -1515,3 +1515,15 @@ def sqrt(cls, series):
     assert Schema.validate(df).equals(  # type: ignore [attr-defined]
         pd.DataFrame({"a": [11.0], "abc": [1.0], "cba": [200.0]})
     )
+
+
+def test_pandera_dtype() -> None:
+    class Schema(pa.DataFrameModel):
+        a: Series[pa.Float]
+        b: Series[pa.Int]
+        c: Series[pa.String]
+
+    df = pd.DataFrame({"a": [1.0], "b": [1], "c": ["1"]})
+    assert Schema.validate(df).equals(  # type: ignore [attr-defined]
+        pd.DataFrame({"a": [1.0], "b": [1], "c": ["1"]})
+    )