Make `remove_key` a pure function

Also, remove duplicate `remove_collected_key` function.
Add doctest to show that it's pure.
Refactor collect functions

- Factor the collect functions out into separate functions rather than
  switching so that they can be more easily extended and indepdendently
  tested.
- Move the collect function tests into doctests. They're small and
  focused so it makes sense for the the tests to be close to the code.