如何重写Python装饰器log_step以控制是否打印被装饰函数的入参

错误原因 #

你当前使用的是无参装饰器的结构，无法直接接收配置参数。使用@log_step(print_params=True)语法时，Python会先执行log_step(print_params=True)，试图将布尔值True作为被装饰函数传入原方法的第一个参数func，触发类型不匹配报错。

修正后的完整代码 #

from functools import wraps
import inspect
import datetime as dt  # 原代码缺失dt的导入，此处补充

def log_step(print_params = False):
    # 外层接收装饰器配置参数，返回实际装饰器
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            tic = dt.datetime.now()
            result = func(*args, **kwargs)
            time_taken = str(dt.datetime.now() - tic)
            print(f"function {func.__name__} took {time_taken}s")
            
            if print_params:
                func_args = inspect.signature(func).bind(*args, **kwargs).arguments
                func_args_str = ", ".join(map("{0[0]} = {0[1]!r}".format, func_args.items()))
                print(f"{func.__name__} arguments: \n==================\n{func_args_str} \n==================\n")        
            
            return result
        return wrapper
    return decorator

def _simple_inside_function(name:str, age:int) -> str:
    """
    打印姓名和年龄的内部工具函数
    """
    age = age + 10
    print(f"Hello {name}, who is {age} years old.")
    
    return age

@log_step(print_params = True)
def wrapper_print_name(config: dict) -> str:
    """
    包装业务函数的外层函数
    """
    A =  _simple_inside_function(**config)
    return A

# 测试用例
if __name__ == "__main__":
    test_config = {"name": "张三", "age": 20}
    wrapper_print_name(test_config)

兼容无括号用法（可选） #

如果需要同时支持@log_step（不填参数默认关闭入参打印）和@log_step(print_params=True)两种写法，可以在外层加参数判断：

def log_step(func=None, print_params = False):
    if func is not None:
        # 处理@log_step无括号的调用场景
        return log_step(print_params=False)(func)
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            tic = dt.datetime.now()
            result = func(*args, **kwargs)
            time_taken = str(dt.datetime.now() - tic)
            print(f"function {func.__name__} took {time_taken}s")
            
            if print_params:
                func_args = inspect.signature(func).bind(*args, **kwargs).arguments
                func_args_str = ", ".join(map("{0[0]} = {0[1]!r}".format, func_args.items()))
                print(f"{func.__name__} arguments: \n==================\n{func_args_str} \n==================\n")        
            
            return result
        return wrapper
    return decorator

内容的提问来源于stack exchange，提问作者Petr