Scala course 10/10 - cleanup code