Fix an infinite loop when using `# fmt: on/off` ... (#3158)

... in the middle of an expression or code block by adding a missing return.

Co-authored-by: Richard Si <63936253+ichard26@users.noreply.github.com>
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>