Little code refactor