Python refactoring