I think that new signature makes sense.  My second comment was about the fact that your diff moves register() from one part of the file to another.  I don't think that's necessary: it should work find in its original location (just below run()) even with the changed signature.