Issue #6 Fixed

The issue has been fixed along with PEP 8 style fixes