I've updated the code. Is this the way to go?