自定义类装饰器实现问题：JobManager无法正常作为装饰器使用

问题描述 #

原本的JobManager作为上下文管理器使用完全正常：

with JobManager(*config) as job:
    <处理逻辑>
# 后续无关处理

但想把它直接当装饰器用在需要完全处于作业上下文的函数上时，出现参数混淆问题，无法正常工作：

@JobManager
def process_job(job, *other_args, **other_kwargs):
    <函数处理逻辑>

相关代码如下：

from contextlib import ContextDecorator, AbstractContextManager, wraps

# 假设存在全局的_jobs变量，以及push_job、pop_job函数
_jobs = []
def push_job(*args, **kwargs):
    return f"job_{len(_jobs)+1}"
def pop_job():
    if _jobs:
        _jobs.pop()

class JobManager(ContextDecorator, AbstractContextManager):
    """同时提供上下文管理器和装饰器功能，在push_job()/pop_job()之间执行代码"""

    def __init__(self, *args, **kwargs):
        self._func, self._args = (
            (args[0], args[1:] if len(args) > 1 else [])
            if args and callable(args[0])
            else (None, args)
        )
        self._kwargs = kwargs

    def __call__(self, func):
        """作为装饰器调用时，返回包裹后的函数"""
    
        @wraps(func)
        def inner(*args, **kwargs):
            with self._recreate_cm() as job_:
                return func(job_, *args, **kwargs)
    
        return inner

    def __enter__(self):
        """上下文管理器进入时，推送作业帧"""
        return (
            push_job(*(self._args), **(self._kwargs))
            if not _jobs else 
            push_job()
            )

    def __exit__(self, exc_type, exc_value, traceback):
        """上下文管理器退出时，弹出作业帧"""
        pop_job()
        return super().__exit__(exc_type, exc_value, traceback)

问题原因 #

1. __init__逻辑混淆：当用@JobManager（无参）装饰函数时，__init__会把被装饰函数当成第一个参数存入self._func，但后续__call__又会再次接收被装饰函数，导致参数处理冲突。
2. _recreate_cm未正确传递参数：ContextDecorator的默认_recreate_cm会调用无参构造，原代码未重写该方法，导致重新创建的实例丢失初始化参数（带参装饰器场景下）。
3. 装饰器参数传递逻辑错误：原__call__中直接复用当前实例创建上下文，但当前实例在装饰器模式下已经绑定了被装饰函数，导致上下文初始化异常。

修复方案 #

重新设计类的初始化和装饰器逻辑，明确区分上下文管理器、无参装饰器、带参装饰器三种场景：

from contextlib import ContextDecorator, AbstractContextManager, wraps

_jobs = []
def push_job(*args, **kwargs):
    return f"job_{len(_jobs)+1}"
def pop_job():
    if _jobs:
        _jobs.pop()

class JobManager(ContextDecorator, AbstractContextManager):
    """同时支持上下文管理器和装饰器的作业管理类，在push_job()和pop_job()之间执行代码"""

    def __init__(self, *args, **kwargs):
        # 只保存push_job需要的参数，不处理被装饰函数
        self.push_args = args
        self.push_kwargs = kwargs

    def __call__(self, func):
        """作为装饰器调用时，返回包裹后的函数"""
        @wraps(func)
        def inner(*args, **kwargs):
            # 每次调用都创建新的上下文管理器实例，确保作业上下文独立
            with self.__class__(*self.push_args, **self.push_kwargs) as job_:
                return func(job_, *args, **kwargs)
        
        return inner

    def __enter__(self):
        """上下文管理器进入时，创建/推送作业"""
        if not _jobs:
            return push_job(*self.push_args, **self.push_kwargs)
        else:
            return push_job()

    def __exit__(self, exc_type, exc_value, traceback):
        """上下文管理器退出时，弹出作业"""
        pop_job()
        return super().__exit__(exc_type, exc_value, traceback)

验证使用场景 #

1. 上下文管理器模式（正常工作）：

with JobManager("config1", "config2") as job:
    print(f"当前作业：{job}")

2. 无参装饰器模式（修复后正常工作）：

@JobManager
def process_job(job, arg1, arg2):
    print(f"作业{job}处理参数：{arg1}, {arg2}")

process_job("a", "b")

3. 带参装饰器模式（新增支持）：

@JobManager("config1", "config2")
def process_job_with_config(job, arg1):
    print(f"带配置的作业{job}处理参数：{arg1}")

process_job_with_config("test")

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