Python decorators

From wikinotes

Decorators let you wrap a function-call in another function.

Decorator Functions

def add_sandwich_suffix(func):
    @functools.wraps(func)
    def addSandwich():
        return func() + ' Sandwich'        
    return addSandwich
@add_sandwich_suffix
def print_ham():
    return 'Ham'

print_ham()
>>> 'Ham Sandwich'

Decorator Classes

class AddSandwichSuffix:
    def __init__(self, function): 
        self.function = function 
  
    def __call__(self): 
        return self.function() + 'Sandwich'
@AddSandwichSuffix
def print_ham():
    return 'Ham'

print_ham()
>>> 'Ham Sandwich'

Decorators with arguments

You can configure your decorator by wrapping it a function.

  • the function returns the decorator
  • the decorator returns the wrapped function

Here is an example I used to mock `input` vs `raw_input` in python 3/2.

def mock_input(return_value):
    def mock_input_decorator(method):
        if sys.version_info[0] < 3:
            patch_path = '__builtin__.raw_input'
        else:
            patch_path = 'builtins.input'

        @mock.patch(patch_path, return_value=return_value)
        @functools.wraps(method)
        def wrapper(message, m_input):
            return method(message, m_input)
        return wrapper

    return mock_input_decorator
class TestSomething
    @mock_input(b'abc')
    def test_something(self, m_input):
        if sys.version_info[0] 3:
            result = raw_input('press enter to continue')
        else:
            result = input('press enter to continue')
        assert result == 'abc'