Modify the code to be a bit more pythonic
Merge pull request #2 from oz123/master

Modify the code to be a bit more pythonic