This looks reasonable. Can you add a test that fails without this change? That would help a lot to prove the correctness of this change and prevent future regressions.