ocwirk
ocwirk

Reputation: 1089

Using python decorator to decorate a part of function

I am trying to mock our git wrapper, so that we can test it. I plan to use mockproc python library which provides the functionality to mock any process name, with a provided script. It works something like this -

self.scripts.append( 'process-name', returncode=0, stdout="output to process" )
with self.scripts:
        run_and_handle_result()

I need to add a decorator layer over this so that I can do some extra things like handle retries. What I want is something like this -

@mockproc('git') # tells that we are mocking git
def test_something(mock_proc):
    mock_proc.set_script("sleep (60)")
    # Run some git command
    mockproc.check_exit_signal()

The problem is I want my decorator to handle the with self.scripts part. So what I want is that the decorator runs the function, setting the process name as git, which is simple. Then run the test function, which adds the script and add with self.script around the git command and then resumes the function.

Is there anyway to do it ? is a decorator bad way to implement it ? This is not a cosmetic requirement. I need this because in some of my commands there is retry logic, for which I need to provide more than one script to mockproc and run multiple times.

Upvotes: 0

Views: 84

Answers (2)

Wyrmwood
Wyrmwood

Reputation: 3589

This is an older question, but I recently encountered a similar situation, where I needed wrap operations, but some were calling other, already wrapped operations, so I didn't want to recursively wrap again. I solved it with closures.

Example wrapper

def wrapper(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except NamedException:
            # do stuff to handle exception
            raise
     
    return decorated

Then I had several functions using the @wrapper


@wrapper
def some_function(stuff):
    # does some stuff

Then I had another function, calling one of those wrapped functions, so to prevent recursively wrapping, while still getting the wrap behavior in the new function, I made a closure.

def some_other_function(args):
    some_function(args)  # already wrapped

    @wrapper
    def _some_other_function()
         # do stuff with args
         return result

    return _some_other_function()

Upvotes: 0

Kijewski
Kijewski

Reputation: 26022

If I understood you correctly, you want to override a named free variable of a function. You can use fun.func_globals[some_name] = some_value. E.g.

def x(a):
    pow2(a)

 x.func_globals['pow2'] = lambda y: y*y

 x(3) == 9

Upvotes: 1

Related Questions