exercise 5.5 in Think Python