have a look to the rebased code, I think you can reuse the change that Andrew made here lately.