Little class refactoring