More on Decorators

• Real-world examples
  • Print the return type
  • Counter
• Decorators and metadata
  • Preserving docstrings when decorating functions
  • Measure decorator overhead
• Decorators that takes arguments
  • Example
  • run_n_times()
  • HTML Generator
• timeout(): a real world example
  • Example
  • Tag your functions
  • Check the return type

Real-world examples

Print the return type

def print_return_type(func):
    # Define wrapper(), the decorated function
    def wrapper(*args, **kwargs):
        # Call the function being decorated
        result = func(*args, **kwargs)
        print('{}() returned type {}'.format(func.__name__, type(result)))
        return result
    # Return the decorated function
    return wrapper

@print_return_type
def foo(value):
    return value

print(foo(42))
print(foo([1, 2, 3]))
print(foo({'a': 42}))

foo() returned type <class 'int'>
42
foo() returned type <class 'list'>
[1, 2, 3]
foo() returned type <class 'dict'>
{'a': 42}

Counter

def counter(func):
    def wrapper(*args, **kwargs):
        wrapper.count += 1
        # Call the function being decorated and return the result
        return func()
    wrapper.count = 0
    # Return the new decorated function
    return wrapper

# Decorate foo() with the counter() decorator
@counter
def foo():
    print('calling foo()')
    
foo()
foo()

print('foo() was called {} times.'.format(foo.count))

calling foo()
calling foo()
foo() was called 2 times.

Decorators and metadata

Preserving docstrings when decorating functions

def add_hello(func):
    def wrapper(*args, **kwargs):
        """Print 'hello' and then call the decorated function."""
        print('hello')
        return func(*args, **kwargs)
    return wrapper

# Decorate print_sum() with the add_hello() decorator
@add_hello
def print_sum(a, b):
    """Adds two numbers and prints the sum"""
    print(a + b)
    
print_sum(10, 20)
print(print_sum.__doc__)

hello
30
Print 'hello' and then call the decorated function.

from functools import wraps

def add_hello(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        """Print 'hello' and then call the decorated function."""
        print('hello')
        return func(*args, **kwargs)
    return wrapper

# Decorate print_sum() with the add_hello() decorator
@add_hello
def print_sum(a, b):
    """Adds two numbers and prints the sum"""
    print(a + b)
    
print_sum(10, 20)
print(print_sum.__doc__)

hello
30
Adds two numbers and prints the sum

Measure decorator overhead

import numpy as np

def check_inputs(a, *args, **kwargs):
    for value in a:
        time.sleep(0.01)
    print('Finished checking inputs')

def check_outputs(a, *args, **kwargs):
    for value in a:
        time.sleep(0.01)
    print('Finished checking outputs')

def check_everything(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        check_inputs(*args, **kwargs)
        result = func(*args, **kwargs)
        check_outputs(result)
        return result
    return wrapper

import time

@check_everything
def duplicate(my_list):
    """Return a new list that repeats the input twice"""
    return my_list + my_list

t_start = time.time()
duplicated_list = duplicate(list(range(50)))
t_end = time.time()
decorated_time = t_end - t_start

t_start = time.time()
# Call the original function instead of the decorated one
duplicated_list = duplicate.__wrapped__(list(range(50)))
t_end = time.time()
undecorated_time = t_end - t_start

print('Decorated time: {:.5f}s'.format(decorated_time))
print('Undecorated time: {:.5f}s'.format(undecorated_time))

Finished checking inputs
Finished checking outputs
Decorated time: 1.51825s
Undecorated time: 0.00025s

Decorators that takes arguments

Example

def run_n_times(n):
    """Define and return a decorator"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            for i in range(n):
                func(*args, **kwargs)
        return wrapper
    return decorator
    
@run_n_times(3)
def print_sum(a, b):
    print(a + b)

print_sum(10, 5)

15
15
15

run_three_times = run_n_times(3)

@run_three_times
def print_sum(a, b):
    print(a + b)
    
print_sum(20, 30)

50
50
50

run_n_times()

def run_n_times(n):
    """Define and return a decorator"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            for i in range(n):
                func(*args, **kwargs)
        return wrapper
    return decorator

@run_n_times(10)
def print_sum(a, b):
    print(a + b)
    
print_sum(15, 20)

35
35
35
35
35
35
35
35
35
35

run_five_times = run_n_times(5)

@run_five_times
def print_sum(a, b):
    print(a + b)
    
print_sum(4, 100)

104
104
104
104
104

HTML Generator

def bold(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        msg = func(*args, **kwargs)
        return '<b>{}</b>'.format(msg)
    return wrapper

def italics(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        msg = func(*args, **kwargs)
        return '<i>{}</i>'.format(msg)
    return wrapper

def html(open_tag, close_tag):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            msg = func(*args, **kwargs)
            return '{}{}{}'.format(open_tag, msg, close_tag)
        # Return the decorated function
        return wrapper
    # Return the decorator
    return decorator

@html('<b>', '</b>')
def hello(name):
    return 'Hello {}!'.format(name)

print(hello('Alice'))

<b>Hello Alice!</b>

@html('<i>', '</i>')
def goodbye(name):
    return 'Goodbye {}.'.format(name)

print(goodbye('Alice'))

<i>Goodbye Alice.</i>

@html('<div>', '</div>')
def hello_goodbye(name):
    return '\n{}\n{}\n'.format(hello(name), goodbye(name))

print(hello_goodbye('Alice'))

<div>
<b>Hello Alice!</b>
<i>Goodbye Alice.</i>
</div>

timeout(): a real world example

Example

import signal
def raise_timeout(*args, **kwargs):
    raise TimeoutError()
    
# When an 'alarm' signal goes off, call raise_timeout()
signal.signal(signalnum=signal.SIGALRM, handler=raise_timeout)

# Set off an alarm in 5 seconds
signal.alarm(5)

# Cancel the alarm
signal.alarm(0)

5

def timeout_in_5s(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        # Set an alarm for 5 seconds
        signal.alarm(5)
        try:
            # Call the decorated func
            return func(*args, **kwargs)
        finally:
            # Cancel alarm
            signal.alarm(0)
    return wrapper

@timeout_in_5s
def foo():
    time.sleep(10)
    print('foo!')

foo()

---------------------------------------------------------------------------
TimeoutError                              Traceback (most recent call last)
<ipython-input-35-c19b6d9633cf> in <module>
----> 1 foo()

<ipython-input-33-ae09c017bfbd> in wrapper(*args, **kwargs)
      6         try:
      7             # Call the decorated func
----> 8             return func(*args, **kwargs)
      9         finally:
     10             # Cancel alarm

<ipython-input-34-79d4e0bed7e7> in foo()
      1 @timeout_in_5s
      2 def foo():
----> 3     time.sleep(10)
      4     print('foo!')

<ipython-input-32-ea109267a28d> in raise_timeout(*args, **kwargs)
      1 import signal
      2 def raise_timeout(*args, **kwargs):
----> 3     raise TimeoutError()
      4 
      5 # When an 'alarm' signal goes off, call raise_timeout()

TimeoutError: 

def timeout(n_seconds):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Set an alarm for n seconds
            signal.alarm(n_seconds)
            try:
                return func(*args, **kwargs)
            finally:
                signal.alarm(0)
        return wrapper
    return decorator

@timeout(20)
def bar():
    time.sleep(10)
    print('bar!')

bar()

bar!

Tag your functions

def tag(*tags):
    # Define a new decorator, named "decorator", to return
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Call the function being decorated and return the result
            return func(*args, **kwargs)
        wrapper.tags = tags
        return wrapper
    # Return the new decorator
    return decorator

@tag('test', 'this is a tag')
def foo():
    pass

print(foo.tags)

('test', 'this is a tag')

Check the return type

def returns_dict(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        assert(type(result) == dict)
        return result
    return wrapper

@returns_dict
def foo(value):
    return value

try:
    print(foo([1, 2, 3]))
except AssertionError:
    print('foo() did not return a dict')

foo() did not return a dict

def returns(return_type):
    # Complete the returns() decorator
    def decorator(func):
        def wrapper(*args, **kwargs):
            result = func(*args, **kwargs)
            assert(type(result) == return_type)
            return result
        return wrapper
    return decorator
  
@returns(dict)
def foo(value):
    return value

try:
    print(foo([1,2,3]))
except AssertionError:
    print('foo() did not return a dict!')

foo() did not return a dict!