import contextlib def set_context(session, context): session.send_session_command("POST", "moz/context", {"context": context}) @contextlib.contextmanager def using_context(session, context): orig_context = session.send_session_command("GET", "moz/context") needs_change = context != orig_context if needs_change: set_context(session, context) try: yield finally: if needs_change: set_context(session, orig_context)