Refactor according to pylint