feat: have set_env return a context manager (#329)

use the Open helper so that `set_env` can be used in with statements,
restoring the original env when exiting the block.

boa/__init__.py

@@ -37,25 +37,20 @@
 
 @contextlib.contextmanager
 def swap_env(new_env):
-    old_env = env
-    try:
-        set_env(new_env)
+    with set_env(new_env):
         yield
-    finally:
-        set_env(old_env)
 
 
-def set_env(new_env):
+def _set_env(new):
     global env
-    env = new_env
+    env = new
+    Env._singleton = new
 
-    Env._singleton = new_env
 
-
-def _env_mgr(new_env):
+def set_env(new_env):
     global env
     get_env = lambda: env  # noqa: E731
-    return Open(get_env, set_env, new_env)
+    return Open(get_env, _set_env, new_env)
 
 
 def fork(
@@ -69,20 +64,20 @@ def fork(
 
     new_env = Env()
     new_env.fork(url=url, block_identifier=block_identifier, deprecated=False, **kwargs)
-    return _env_mgr(new_env)
+    return set_env(new_env)
 
 
 def set_browser_env(address=None):
     """Set the environment to use the browser's network in Jupyter/Colab"""
     # import locally because jupyter is generally not installed
     from boa.integrations.jupyter import BrowserEnv
 
-    return _env_mgr(BrowserEnv(address))
+    return set_env(BrowserEnv(address))
 
 
 def set_network_env(url):
     """Set the environment to use a custom network URL"""
-    return _env_mgr(NetworkEnv.from_url(url))
+    return set_env(NetworkEnv.from_url(url))
 
 
 def set_etherscan(*args, **kwargs):

tests/unitary/test_boa.py

@@ -4,7 +4,7 @@
 def test_env_mgr_noctx():
     s = boa.env
     t = boa.Env()
-    boa._env_mgr(t)
+    boa.set_env(t)
     assert boa.env is not s
     assert boa.env is t
 
@@ -13,7 +13,7 @@ def test_env_mgr_with_ctx():
     s = boa.env
     t = boa.Env()
 
-    with boa._env_mgr(t):
+    with boa.set_env(t):
         assert boa.env is not s
         assert boa.env is t