Refactored state_context method to use contextmanager decorator

And more docstrings