Adding __enter__ and __exit__ definitions