I thought I understood this, but clearly I don't.

I have a context manager that wraps suites of code inside some metadata handling. The details of that don't matter, suffice to say the code needs to be between a push_job() and pop_job(). That all works, so I can say:

with JobManager(*config) as job:
    <Processing Goes Here>
<Follow up processing, not connected to the job>

I thought it would be neat if, instead of inserting the with statement when an entire function needed to be in its own job context, I could just use the same class as a decorator, and say:

@JobManager
def process_job(job, *other_args, *other_kwargs):

But it's not working as a decorator, and I can't see why. I get confusion of arguments between the decorator and the wrapped function.

Here's the code:

class JobManager(ContextDecorator, AbstractContextManager):
    """Convenience class that provides both a context manager to surround code
    suites needing a new job frame (using push_job()/pop_job()), and a
    decorator to wrap functions in the current frame automagically
    """

    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):
        """Called when using JobManager as a decorator. Returns a new frame."""
    
        @wraps(func)
        def inner(*args, **kwargs):
            with self._recreate_cm() as job_:
                return func(job_, *args, **kwargs)
    
        return inner

    def __enter__(self):
        """Called when using JobManager as a context manager. Pushes a new
        frame or creates the initial frame from the arguments.
        """
        return (
            push_job(*(self._args), **(self._kwargs))
            if not _jobs else 
            push_job()
            )

    def __exit__(self, exc_type, exc_value, traceback):
        """Called when using JobManager as a context manager."""
        pop_job()
        return super().__exit__(exc_type, exc_value, traceback)

CodePudding user response：

Doing this:

@some_decorator
def func(...):
    ...

Is same as doing this:

def func(...):
    ...

func = some_decorator(func)

So what you are doing is basically process_job = JobManager(process_job). In other words, you just created an instance of JobManager. This clearly cannot work, as now process_job() will invoke the __call__() method of the instance, which will simply return the inner function and do nothing.

You need to put your "decorating" logic inside your __call__() method, and decorate with @JobManager() instead (notice the added () there). This is because __init__() should not return anything. Otherwise, you can avoid the inner wrapper inside your __call__() method and do the job directly there:

    def __call__(self, *args, **kwargs):
        with self._recreate_cm() as job_:
            return self._func(job_, *args, **kwargs)

The only "issue" here is that you lose the benefits of @wraps. Also, be careful, because then you will need some sanity checks to make sure that calling a normal JobManager instance does not raise exceptions because there is no associated function to run.

In any case, I don't see the need for all of this complex fuss. Just define a normal decorator that is a function and leverage the closure of the function to do whatever you need.

def managed(func):
    # You could also create the JobManager here if you want
    # e.g. job = JobManager(), and then use it in the wrapper,
    # but that'd need manual enter/exit.

    @wraps(func)
    def wrapper(*a, **kwa):
        with JobManager() as job:
            return func(job, *a, **kwa)

    return wrapper

@managed
def process_job(job, ...):
    ...

CodePudding user response：

Marco's answer addresses your issue nicely. I'd just like to share my version of the decorator and context manager hybrid from my open source library pip install generallibrary

Definition for DecoContext is here

It would let you write your use case very cleanly like this:

from generallibrary import DecoContext

class JobManager(DecoContext):
    def before(self):
        push_job()

    def after(self):
        pop_job()

Lots of test cases here

Defining the __get__ method to return a bound __call__ method is vital to be able to decorate methods