Edits after running with pylint