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]>
unionai-oss · Nov 13, 2024 · 9667234 · 9667234
1 parent a0ac2a1
commit 9667234
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/pandera/dtypes.py b/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]]:

diff --git a/tests/core/test_model.py b/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"]})
+    )