diff --git a/modin/logging/logger_decorator.py b/modin/logging/logger_decorator.py index 301fb02562b..d8e457479fc 100644 --- a/modin/logging/logger_decorator.py +++ b/modin/logging/logger_decorator.py @@ -17,9 +17,11 @@ ``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 @@ -27,8 +29,10 @@ _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!"