feat: add overload.

Signed-off-by: yassun7010 <[email protected]>
unionai-oss · Sep 30, 2024 · 62073f4 · 62073f4
1 parent 0207fba
commit 62073f4
Showing 1 changed file with 14 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,18 @@ def __hash__(self) -> int:
 _DataTypeClass = Type[_Dtype]
 
 
+@overload
+def immutable(
+    pandera_dtype_cls: _DataTypeClass, **dataclass_kwargs: Any
+) -> _DataTypeClass:
+    ...
+
+@overload
+def immutable(
+    pandera_dtype_cls: None = None, **dataclass_kwargs: Any
+) -> Callable[[_DataTypeClass], _DataTypeClass]:
+    ...
+
 def immutable(
     pandera_dtype_cls: Optional[_DataTypeClass] = None, **dataclass_kwargs: Any
 ) -> Union[_DataTypeClass, Callable[[_DataTypeClass], _DataTypeClass]]: