Minor refactor of repo

Use return over if/else branching in repo.