Minor class refactoring