This looks great, but I think it'd be great if you submitted this (if you haven't already) as an issue to Jetbrains! They may integrate a similar fix and that'd be much more discoverable