Extend API to allow specifying the exception class

Adds an `on_exception` argument to specify the class of exception to catch. This provides a cleaner API to catching only specific exception classes.

Opening this in part as a POC and to start discussion. I think it is a much nicer API for this use case and also solves the issue of catching `BaseException` by default.

See #17