Cleaned up few lines of code

I am not very comfortable with the usage of `six`, which is not a built in module.