Refactor for PR