small python refactor