Small improvements

Very small :) 

Do pep8 restyling and use more pythonic solutions.