Yes! This is the right fix IMO.
I wanted to check that this wasn't a 'hack'

But is correct - python's tokenizer parses '@' as an operator (tokenize.OP) because it's actually a matrix multiplication operator, as well as being used for function decoration.

So '@' used to end up in the buffer that was flushed in the case of a context_exit triggered specifically after a run block. This left black failing to parse a lone @. (Note after something like input:, buffer used to still contain a lone '@' but buffer flushing would not be triggered as context_exit is not called). Now we force '@' not to be present in the flushed buffer when calling context_exit (i.e. exiting a run block). We do return a Status containing @ though, but that's OK as it won't be recognised as a keyword, and we'll just continue parsing the snakefile.

Wrote this as a documentation if we ever come back to it.
Nice fix! 👍