Set chain ID via BrowserSigner + #137

Makefile

@@ -3,9 +3,7 @@
 all: lint build
 
 lint:
-	black -C -t py310 boa/ tests/
-	isort boa/ tests/
-	flake8 boa/ tests/
+	pre-commit run --all-files
 	mypy --install-types --non-interactive --follow-imports=silent --ignore-missing-imports --implicit-optional -p boa
 
 build:

boa/__init__.py

@@ -47,11 +47,9 @@ def 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 BrowserRPC, BrowserSigner
+    from boa.integrations.jupyter import BrowserEnv
 
-    env = NetworkEnv(rpc=BrowserRPC())
-    env.set_eoa(BrowserSigner(address))
-    set_env(env)
+    set_env(BrowserEnv(address))
 
 
 def set_network_env(url):

boa/contracts/vyper/vyper_contract.py

@@ -33,16 +33,27 @@
 from vyper.utils import method_id
 
 from boa import BoaError
-from boa.contracts.base_evm_contract import StackTrace, _BaseEVMContract, _handle_child_trace
-from boa.contracts.vyper.ast_utils import ast_map_of, get_fn_ancestor_from_node, reason_at
+from boa.contracts.base_evm_contract import (
+    StackTrace,
+    _BaseEVMContract,
+    _handle_child_trace,
+)
+from boa.contracts.vyper.ast_utils import (
+    ast_map_of,
+    get_fn_ancestor_from_node,
+    reason_at,
+)
 from boa.contracts.vyper.compiler_utils import (
     _METHOD_ID_VAR,
     anchor_compiler_settings,
     compile_vyper_function,
     generate_bytecode_for_arbitrary_stmt,
     generate_bytecode_for_internal_fn,
 )
-from boa.contracts.vyper.decoder_utils import ByteAddressableStorage, decode_vyper_object
+from boa.contracts.vyper.decoder_utils import (
+    ByteAddressableStorage,
+    decode_vyper_object,
+)
 from boa.contracts.vyper.event import Event, RawEvent
 from boa.contracts.vyper.ir_executor import executor_from_ir
 from boa.environment import Env

boa/environment.py

@@ -399,14 +399,9 @@ def set_random_seed(self, seed=None):
     def get_gas_price(self):
         return self._gas_price or 0
 
-    def _set_account_db_class(self, account_db_class: type):
-        self.vm.__class__._state_class.account_db_class = account_db_class
-
     def _init_vm(self, reset_traces=True, account_db_class=AccountDB):
         self.vm = self.chain.get_vm()
-
-        self._set_account_db_class(account_db_class)
-
+        self.vm.__class__._state_class.account_db_class = account_db_class
         self.vm.patch = VMPatcher(self.vm)
 
         c = type(

boa/integrations/jupyter/__init__.py

@@ -1,4 +1,4 @@
-from boa.integrations.jupyter.browser import BrowserRPC, BrowserSigner
+from boa.integrations.jupyter.browser import BrowserEnv, BrowserRPC, BrowserSigner
 from boa.integrations.jupyter.constants import PLUGIN_NAME
 from boa.integrations.jupyter.handlers import setup_handlers
 
@@ -19,6 +19,7 @@ def load_jupyter_server_extension(server_app):
 __all__ = [
     BrowserSigner,
     BrowserRPC,
+    BrowserEnv,
     load_jupyter_server_extension,
     _load_jupyter_server_extension,
 ]

boa/integrations/jupyter/browser.py

@@ -13,9 +13,7 @@
 import nest_asyncio
 from IPython.display import Javascript, display
 
-from boa.rpc import RPC, RPCError
-
-from .constants import (
+from boa.integrations.jupyter.constants import (
     ADDRESS_TIMEOUT_MESSAGE,
     CALLBACK_TOKEN_BYTES,
     CALLBACK_TOKEN_TIMEOUT,
@@ -25,7 +23,12 @@
     SHARED_MEMORY_LENGTH,
     TRANSACTION_TIMEOUT_MESSAGE,
 )
-from .utils import convert_frontend_dict, install_jupyter_javascript_triggers
+from boa.integrations.jupyter.utils import (
+    convert_frontend_dict,
+    install_jupyter_javascript_triggers,
+)
+from boa.network import NetworkEnv
+from boa.rpc import RPC, RPCError
 
 try:
     from google.colab.output import eval_js as colab_eval_js
@@ -121,6 +124,31 @@ def wait_for_tx_receipt(self, tx_hash, timeout: float, poll_latency=1):
         )
 
 
+class BrowserEnv(NetworkEnv):
+    """
+    A NetworkEnv object that uses the BrowserSigner and BrowserRPC classes.
+    """
+
+    def __init__(self, address=None):
+        super().__init__(rpc=BrowserRPC())
+        self.signer = BrowserSigner(address)
+        self.set_eoa(self.signer)
+
+    def get_chain_id(self):
+        return _javascript_call(
+            "rpc", "eth_chainId", timeout_message=RPC_TIMEOUT_MESSAGE
+        )
+
+    def set_chain_id(self, chain_id):
+        _javascript_call(
+            "rpc",
+            "wallet_switchEthereumChain",
+            [{"chainId": chain_id}],
+            timeout_message=RPC_TIMEOUT_MESSAGE,
+        )
+        self._reset_fork()
+
+
 def _javascript_call(js_func: str, *args, timeout_message: str) -> Any:
     """
     This function attempts to call a Javascript function in the browser and then

boa/interpret.py

@@ -9,7 +9,11 @@
 
 from boa.contracts.abi.abi_contract import ABIContractFactory
 from boa.contracts.vyper.compiler_utils import anchor_compiler_settings
-from boa.contracts.vyper.vyper_contract import VyperBlueprint, VyperContract, VyperDeployer
+from boa.contracts.vyper.vyper_contract import (
+    VyperBlueprint,
+    VyperContract,
+    VyperDeployer,
+)
 from boa.explorer import fetch_abi_from_etherscan
 from boa.util.abi import Address
 from boa.util.disk_cache import DiskCache

boa/network.py

@@ -9,7 +9,16 @@
 from requests.exceptions import HTTPError
 
 from boa.environment import Env
-from boa.rpc import RPC, EthereumRPC, RPCError, fixup_dict, to_bytes, to_hex, to_int, trim_dict
+from boa.rpc import (
+    RPC,
+    EthereumRPC,
+    RPCError,
+    fixup_dict,
+    to_bytes,
+    to_hex,
+    to_int,
+    trim_dict,
+)
 from boa.util.abi import Address
 
 

tests/unitary/jupyter/test_browser.py

@@ -76,7 +76,7 @@ def find_response(mock_calls, func_to_body_dict):
 
 
 @pytest.fixture()
-def mock_callback(mocked_token, browser, display_mock):
+def mock_callback(mocked_token, browser, display_mock, mock_inject_javascript):
     """Returns a function that allows mocking the result of the frontend callback."""
 
     with mock.patch(
@@ -220,6 +220,19 @@ def test_browser_loads_signer(
     mock_inject_javascript.assert_called()
 
 
+def test_browser_chain_id(token, env, display_mock, mock_callback):
+    mock_callback("eth_chainId", "0x123")
+    assert env.get_chain_id() == "0x123"
+    mock_callback("wallet_switchEthereumChain")
+    env.set_chain_id("0x456")
+    assert display_mock.call_count == 7
+    (js,), _ = display_mock.call_args_list[-2]
+    assert (
+        f'rpc("{token}", "wallet_switchEthereumChain", [{{"chainId": "0x456"}}])'
+        in js.data
+    )
+
+
 def test_browser_rpc(
     token,
     browser,
@@ -240,14 +253,7 @@ def test_browser_rpc(
 
 
 def test_browser_rpc_error(
-    token,
-    browser,
-    display_mock,
-    mock_callback,
-    mock_inject_javascript,
-    account,
-    mock_fork,
-    env,
+    token, browser, display_mock, mock_callback, account, mock_fork, env
 ):
     rpc_error = {"code": -32000, "message": "Reverted"}
     mock_callback(
@@ -259,14 +265,7 @@ def test_browser_rpc_error(
 
 
 def test_browser_rpc_server_error(
-    token,
-    browser,
-    display_mock,
-    mock_callback,
-    mock_inject_javascript,
-    account,
-    mock_fork,
-    env,
+    token, browser, display_mock, mock_callback, account, mock_fork, env
 ):
     error = {
         "code": "UNKNOWN_ERROR",
@@ -279,13 +278,7 @@ def test_browser_rpc_server_error(
 
 
 def test_browser_js_error(
-    token,
-    browser,
-    display_mock,
-    mock_callback,
-    mock_inject_javascript,
-    account,
-    mock_fork,
+    token, browser, display_mock, mock_callback, account, mock_fork
 ):
     mock_callback("loadSigner", error={"message": "custom message", "stack": ""})
     with pytest.raises(browser.RPCError) as exc_info: