FIX-#7278: Make sure enable_logging decorator preserve type hints (#…

…7279)

Signed-off-by: Anatoly Myachev <[email protected]>

modin/logging/logger_decorator.py

@@ -17,18 +17,22 @@
 ``enable_logging`` is used for decorating individual Modin functions or classes.
 """
 
+from __future__ import annotations
+
 from functools import wraps
 from types import FunctionType, MethodType
-from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, overload
 
 from modin.config import LogMode
 
 from .config import LogLevel, get_logger
 
 _MODIN_LOGGER_NOWRAP = "__modin_logging_nowrap__"
 
+Fn = TypeVar("Fn", bound=Any)
+
 
-def disable_logging(func: Callable) -> Callable:
+def disable_logging(func: Callable) -> Any:
     """
     Disable logging of one particular function. Useful for decorated classes.
 
@@ -46,11 +50,26 @@ def disable_logging(func: Callable) -> Callable:
     return func
 
 
+@overload
+def enable_logging(modin_layer: Fn) -> Fn:
+    # This helps preserve typings when the decorator is used without parentheses
+    pass
+
+
+@overload
+def enable_logging(
+    modin_layer: str = "PANDAS-API",
+    name: Optional[str] = None,
+    log_level: LogLevel = LogLevel.INFO,
+) -> Callable[[Fn], Fn]:
+    pass
+
+
 def enable_logging(
-    modin_layer: Union[str, Callable, Type] = "PANDAS-API",
+    modin_layer: str | Fn = "PANDAS-API",
     name: Optional[str] = None,
     log_level: LogLevel = LogLevel.INFO,
-) -> Callable:
+) -> Callable[[Fn], Fn] | Fn:
     """
     Log Decorator used on specific Modin functions or classes.
 
@@ -76,7 +95,7 @@ def enable_logging(
         # def func()
         return enable_logging()(modin_layer)
 
-    def decorator(obj: Any) -> Any:
+    def decorator(obj: Fn) -> Fn:
         """Decorate function or class to add logs to Modin API function(s)."""
         if isinstance(obj, type):
             seen: Dict[Any, Any] = {}
@@ -94,11 +113,11 @@ def decorator(obj: Any) -> Any:
                         )(attr_value)
 
                     setattr(obj, attr_name, wrapped)
-            return obj
+            return obj  # type: ignore [return-value]
         elif isinstance(obj, classmethod):
-            return classmethod(decorator(obj.__func__))
+            return classmethod(decorator(obj.__func__))  # type: ignore [return-value, arg-type]
         elif isinstance(obj, staticmethod):
-            return staticmethod(decorator(obj.__func__))
+            return staticmethod(decorator(obj.__func__))  # type: ignore [return-value, arg-type]
 
         assert isinstance(modin_layer, str), "modin_layer is somehow not a string!"