Refactor to make some if statements one-liners