Yea if you want to refactor it then that's fine. If the refactor can wait, I can do it in a separate pull