Refactor code

[amitness]

• Use PEP 8 best practices
• Remove redundant code
• Use standard library function instead of reinventing