Decorators and context managers

Materials

https://scipy-lectures.github.io/advanced/advanced_python/index.html

Public etherpad instance

Use this to post questions, links, comments, and sometimes solutions to exercises.

https://beta.etherpad.org/p/aspp

Exercises

Exercise 1: Wrapping a function to do something on every invocation

Write a decorator which prints the arguments and the return value of the wrapped function. 

@logger         
def g(x):
    return 2 * x
 
@logger
def f(x, y):
    return g(x) + g(y)
>>> f(5, 6)
f is called with args [5, 6] kwargs {}
g is called with args [5] kwargs {}
g returns 10
g is called with args [6] kwargs {}
g returns 12
f returns 22

Exercise 2: A decorator which times function execution

Write 'timeit' decorator which prints how long a function took to execute. 

import time
 
@timeit
def loooong():
    time.sleep(5)
    return 'ans'
>>> looong()
looong took 5.0323423s
ans

Exercise 3: A cache

Write a decorator which caches the results of some function. Store the results of every invocation. Every time the function is called with the same arguments, simply retrieve the value from storage. Otherwise call the function. 

>>> @cached
>>> def f(x): print('here')

>>> f(3)
here
>>> f(3)
>>>

Test with: 

def factorial(n):
    if n <= 1:
        return 1
    else:
        return n * factorial(n-1)

Exercise 3a: Keeping state in decorators

Write a decorator which prints a warning the first time a given function is executed. This is a modification of deprecate() from previous exercise. 

@deprecate('do not use')
def f():
    pass
>>> f()
f is deprecated, do not use
>>> f()
>>> f()

The trick is how to store the state!

Exercise 4: Returning a list of results

When a function returns a list of results, we might need to gather those results in a list: 

def lucky_numbers(n):
    ans = []
    for i in range(n):
        if i % 7 != 0:
            continue
        if sum(int(digit) for digit in str(i)) % 3 != 0:
            continue
        ans.append(i)
    return ans

This looks much nicer when written as a generator. First convert lucky_numbers to be a generator.

Later, write listize decorator which gathers the results from a generator and returns a list and use it to wrap the new lucky_numbers().

Alternatively, write arrayize decorator which return the results in a numpy array.

Exercise 5: A context manager to time execution

Before we wrote a decorator which would print how long a function took to execute. Now write a context manager which does the same thing. 

>>> with logtime_cm():
...     time.sleep(3)
Execution took 3.00001s

Exercise 6: A matplotlib context manager

This is synthesized from a real program that I use to analyze results and create graphs. Matplotlib figures can be plotted on screen, and they can also be saved to file with figure.savefig().

Write a context manager which gives you a matplotlib figure object, and either saves the plot to a file or pops it up on screen, depending on a global parameter SAVEFIGS (in a real program this parameter would be settable by a commandline option). 

with save_or_plot('name') as f:
    ax = f.gca()
    ax.plot([0, 3, 2, 5])
    ax.set_xlabel('x')
    ax.set_ylabel('y')

Exercise 7: Checking exception raising

This example comes from unit testing. We want to make sure that we raise the right exceptions on errors.

Write a cm 'assert_raises' that

  1. checks that an exception was raised
  2. checks that the exception is of the right type
>>> with assert_raises(ZeroDivisionError):
...     1 / 0

>>> with assert_raises(ZeroDivisionError):
...     0[0] / 0
Traceback (most recent call last):
  ...
AssertionError: expected ZeroDivisionError not AttributeError

>>> with assert_raises(ZeroDivisionError):
...     0 / 1
Traceback (most recent call last):
  ...
AssertionError: expected ZeroDivisionError exception

Exercise 8: Context manager to limit computation time

The OS call alarm can be used to interrupt a process: 

import signal
import time
 
def _handler(signum, frame):
    print('_handler called for signal', signum)
 
oldhandler = signal.signal(signal.SIGALRM, _handler)
signal.alarm(3)
time.sleep(100)

Write a context manager which limits the execution time to the given number of seconds: 

>>> with timelimit(5):
...     looong_computation()
RuntimeError         Traceback (most recent call last)
...
RuntimeError: over the deadline

Notes: signal.signal returns the previous handler that was installed. It should be restored after our context manager is done.

Some docs on the web: