Fixes #9

Do we want this? Commented out other lines, we can put them
back in if we prefer the previous functionality