fix issue #9

fix issue #9, also did a pep8 format