How about moving this logic to the `cli` module? The code would be more lightweight, and we could add unit tests for it more easily.