Minor refactors to clean up the screen module