Add support for `with` statement