Simplified the code
PEP8