#!/usr/bin/env python3
"""
进度装饰器
"""

def with_progress(total_steps: int = 100, task_name: str = None):
    """
    为函数添加进度报告的装饰器
    
    Args:
        total_steps: 总步数
        task_name: 任务名称
    """
    def decorator(func):
        def wrapper(self, *args, **kwargs):
            name = task_name or func.__name__
            
            if hasattr(self, 'create_task'):
                with self.create_task(name, total_steps) as task:
                    # 将task对象传递给函数
                    return func(self, task, *args, **kwargs)
            else:
                # 如果不是ProgressJSONRPCCommander，直接执行
                return func(self, *args, **kwargs)
        
        return wrapper
    return decorator
