Add typehint for deprecated and experimental

optuna/_deprecated.py

@@ -4,13 +4,21 @@
 from typing import Any
 from typing import Callable
 from typing import Optional
+from typing import overload
+from typing import Type
+from typing import TypeVar
+from typing import Union
 import warnings
 
 from packaging import version
+from typing_extensions import ParamSpec
 
 from optuna._experimental import _get_docstring_indent
 from optuna._experimental import _validate_version
 
+FT = TypeVar("FT")
+FP = ParamSpec("FP")
+CT = TypeVar("CT")
 
 _DEPRECATION_NOTE_TEMPLATE = """
 
@@ -46,7 +54,7 @@ def deprecated(
     removed_version: str,
     name: Optional[str] = None,
     text: Optional[str] = None,
-) -> Any:
+) -> Union[Callable[[CT], CT], Callable[FP, FT]]:
     """Decorate class or function as deprecated.
 
     Args:
@@ -75,10 +83,18 @@ def deprecated(
     _validate_version(removed_version)
     _validate_two_version(deprecated_version, removed_version)
 
+    @overload
+    def _deprecated_wrapper(f: Callable[FP, FT]) -> Callable[FP, FT]:
+        ...
+
+    @overload
+    def _deprecated_wrapper(f: CT) -> CT:
+        ...
+
     def _deprecated_wrapper(f: Any) -> Any:
         # f is either func or class.
 
-        def _deprecated_func(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
+        def _deprecated_func(func: Callable[FP, FT]) -> Callable[FP, FT]:
             """Decorates a function as deprecated.
 
             This decorator is supposed to be applied to the deprecated function.
@@ -106,11 +122,11 @@ def new_func(*args: Any, **kwargs: Any) -> Any:
                     message += " " + text
                 warnings.warn(message, FutureWarning, stacklevel=2)
 
-                return func(*args, **kwargs)  # type: ignore
+                return func(*args, **kwargs)
 
             return new_func
 
-        def _deprecated_class(cls: Any) -> Any:
+        def _deprecated_class(cls: Type[CT]) -> Type[CT]:
             """Decorates a class as deprecated.
 
             This decorator is supposed to be applied to the deprecated class.
@@ -134,7 +150,7 @@ def wrapped_init(self, *args, **kwargs) -> None:  # type: ignore
 
                 _original_init(self, *args, **kwargs)
 
-            cls.__init__ = wrapped_init
+            setattr(cls, "__init__", wrapped_init)
 
             if cls.__doc__ is None:
                 cls.__doc__ = ""

optuna/_experimental.py

@@ -4,11 +4,21 @@
 from typing import Any
 from typing import Callable
 from typing import Optional
+from typing import overload
+from typing import Type
+from typing import TypeVar
+from typing import Union
 import warnings
 
+from typing_extensions import ParamSpec
+
 from optuna.exceptions import ExperimentalWarning
 
 
+FT = TypeVar("FT")
+FP = ParamSpec("FP")
+CT = TypeVar("CT")
+
 _EXPERIMENTAL_NOTE_TEMPLATE = """
 
 .. note::
@@ -31,7 +41,10 @@ def _get_docstring_indent(docstring: str) -> str:
     return docstring.split("\n")[-1] if "\n" in docstring else ""
 
 
-def experimental(version: str, name: Optional[str] = None) -> Any:
+def experimental(
+    version: str,
+    name: Optional[str] = None,
+) -> Union[Callable[[CT], CT], Callable[FP, FT]]:
     """Decorate class or function as experimental.
 
     Args:
@@ -41,10 +54,18 @@ def experimental(version: str, name: Optional[str] = None) -> Any:
 
     _validate_version(version)
 
+    @overload
+    def _experimental_wrapper(f: Callable[FP, FT]) -> Callable[FP, FT]:
+        ...
+
+    @overload
+    def _experimental_wrapper(f: CT) -> CT:
+        ...
+
     def _experimental_wrapper(f: Any) -> Any:
         # f is either func or class.
 
-        def _experimental_func(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
+        def _experimental_func(func: Callable[FP, FT]) -> Callable[FP, FT]:
 
             if func.__doc__ is None:
                 func.__doc__ = ""
@@ -65,19 +86,19 @@ def new_func(*args: Any, **kwargs: Any) -> Any:
                     stacklevel=2,
                 )
 
-                return func(*args, **kwargs)  # type: ignore
+                return func(*args, **kwargs)
 
             return new_func
 
-        def _experimental_class(cls: Any) -> Any:
+        def _experimental_class(cls: Type[CT]) -> Type[CT]:
             """Decorates a class as experimental.
 
             This decorator is supposed to be applied to the experimental class.
             """
             _original_init = cls.__init__
 
             @functools.wraps(_original_init)
-            def wrapped_init(self, *args, **kwargs) -> None:  # type: ignore
+            def wrapped_init(self, *args: Any, **kwargs: Any) -> None:  # type: ignore
                 warnings.warn(
                     "{} is experimental (supported from v{}). "
                     "The interface can change in the future.".format(
@@ -89,7 +110,7 @@ def wrapped_init(self, *args, **kwargs) -> None:  # type: ignore
 
                 _original_init(self, *args, **kwargs)
 
-            cls.__init__ = wrapped_init
+            setattr(cls, "__init__", wrapped_init)
 
             if cls.__doc__ is None:
                 cls.__doc__ = ""