Docstring fix

Quick edit to docstring and some PEP8 cleaning