Refactor to Decorator pattern.