More style refactoring.
Refactoring
Added some type annotations in a couple of places to make my IDE a bit happier.