refactor: Simplify the jupyter plugin

boa/integrations/jupyter/browser.py

@@ -4,10 +4,12 @@
 """
 import json
 import logging
-from asyncio import get_running_loop, sleep
+import os
+from asyncio import get_event_loop, sleep
 from itertools import chain
 from multiprocessing.shared_memory import SharedMemory
 from os import urandom
+from os.path import dirname, join, realpath
 from typing import Any
 
 import nest_asyncio
@@ -22,10 +24,6 @@
     SHARED_MEMORY_LENGTH,
     TRANSACTION_TIMEOUT_MESSAGE,
 )
-from boa.integrations.jupyter.utils import (
-    convert_frontend_dict,
-    install_jupyter_javascript_triggers,
-)
 from boa.network import NetworkEnv
 from boa.rpc import RPC, RPCError
 from boa.util.abi import Address
@@ -39,26 +37,85 @@
 nest_asyncio.apply()
 
 
-if not colab_eval_js:
-    # colab creates a new iframe for every call, we need to re-inject it every time
-    # for jupyterlab we only need to do it once
-    install_jupyter_javascript_triggers()
+def _install_javascript_triggers():
+    """Run the ethers and titanoboa_jupyterlab Javascript snippets in the browser."""
+    cur_dir = dirname(realpath(__file__))
+    with open(join(cur_dir, "jupyter.js")) as f:
+        js = f.read()
+
+    prefix = os.getenv("JUPYTERHUB_SERVICE_PREFIX", "..")
+    js = js.replace("$$JUPYTERHUB_SERVICE_PREFIX", prefix)
+    js = js.replace("$$BOA_DEBUG_MODE", json.dumps(BrowserRPC._debug_mode))
+
+    display(Javascript(js))
+
+
+class BrowserRPC(RPC):
+    """
+    An RPC object that sends requests to the browser via Javascript.
+    """
+
+    _debug_mode = False
+
+    def __init__(self):
+        if not colab_eval_js:
+            # colab creates a new iframe for every call, we need to re-inject it every time
+            # for jupyterlab we only need to do it once
+            _install_javascript_triggers()
+
+    @property
+    def identifier(self) -> str:
+        return type(self).__name__  # every instance does the same
+
+    @property
+    def name(self):
+        return self.identifier
+
+    def fetch(
+        self, method: str, params: Any, timeout_message=RPC_TIMEOUT_MESSAGE
+    ) -> Any:
+        return _javascript_call("rpc", method, params, timeout_message=timeout_message)
+
+    def fetch_multi(
+        self, payloads: list[tuple[str, Any]], timeout_message=RPC_TIMEOUT_MESSAGE
+    ) -> list[Any]:
+        return _javascript_call("multiRpc", payloads, timeout_message=timeout_message)
+
+    def wait_for_tx_receipt(self, tx_hash, timeout: float, poll_latency=1):
+        # we do the polling in the browser to avoid too many callbacks
+        # each callback generates currently 10px empty space in the frontend
+        timeout_ms, pool_latency_ms = timeout * 1000, poll_latency * 1000
+        return _javascript_call(
+            "waitForTransactionReceipt",
+            tx_hash,
+            timeout_ms,
+            pool_latency_ms,
+            timeout_message=RPC_TIMEOUT_MESSAGE,
+        )
 
 
 class BrowserSigner:
     """
     A BrowserSigner is a class that can be used to sign transactions in IPython/JupyterLab.
     """
 
-    def __init__(self, address=None):
+    def __init__(self, address=None, rpc=None):
         """
         Create a BrowserSigner instance.
         :param address: The account address. If not provided, it will be requested from the browser.
         """
+        if rpc is None:
+            rpc = BrowserRPC()  # note: the browser window is global anyway
+        self._rpc = rpc
         address = getattr(address, "address", address)
-        address = _javascript_call(
-            "loadSigner", address, timeout_message=ADDRESS_TIMEOUT_MESSAGE
-        )
+        accounts = self._rpc.fetch("eth_requestAccounts", [], ADDRESS_TIMEOUT_MESSAGE)
+
+        if address is None and len(accounts) > 0:
+            address = accounts[0]
+
+        if address not in accounts:
+            raise ValueError(f"Address {address} is not available in the browser")
+
         self.address = Address(address)
 
     def send_transaction(self, tx_data: dict) -> dict:
@@ -69,60 +126,21 @@ def send_transaction(self, tx_data: dict) -> dict:
         :param tx_data: The transaction data to sign.
         :return: The signed transaction data.
         """
-        sign_data = _javascript_call(
-            "sendTransaction", tx_data, timeout_message=TRANSACTION_TIMEOUT_MESSAGE
+        hash = self._rpc.fetch(
+            "eth_sendTransaction", [tx_data], TRANSACTION_TIMEOUT_MESSAGE
         )
-        return convert_frontend_dict(sign_data)
+        return {"hash": hash}
 
     def sign_typed_data(self, full_message: dict[str, Any]) -> str:
         """
         Sign typed data value with types data structure for domain using the EIP-712 specification.
         :param full_message: The full message to sign.
         :return: The signature.
         """
-        return _javascript_call(
-            "rpc",
+        return self._rpc.fetch(
             "eth_signTypedData_v4",
             [self.address, full_message],
-            timeout_message=TRANSACTION_TIMEOUT_MESSAGE,
-        )
-
-
-class BrowserRPC(RPC):
-    """
-    An RPC object that sends requests to the browser via Javascript.
-    """
-
-    _debug_mode = False
-
-    @property
-    def identifier(self) -> str:
-        return type(self).__name__  # every instance does the same
-
-    @property
-    def name(self):
-        return self.identifier
-
-    def fetch(self, method: str, params: Any) -> Any:
-        return _javascript_call(
-            "rpc", method, params, timeout_message=RPC_TIMEOUT_MESSAGE
-        )
-
-    def fetch_multi(self, payloads: list[tuple[str, Any]]) -> list[Any]:
-        return _javascript_call(
-            "multiRpc", payloads, timeout_message=RPC_TIMEOUT_MESSAGE
-        )
-
-    def wait_for_tx_receipt(self, tx_hash, timeout: float, poll_latency=1):
-        # we do the polling in the browser to avoid too many callbacks
-        # each callback generates currently 10px empty space in the frontend
-        timeout_ms, pool_latency_ms = timeout * 1000, poll_latency * 1000
-        return _javascript_call(
-            "waitForTransactionReceipt",
-            tx_hash,
-            timeout_ms,
-            pool_latency_ms,
-            timeout_message=RPC_TIMEOUT_MESSAGE,
+            TRANSACTION_TIMEOUT_MESSAGE,
         )
 
 
@@ -131,11 +149,17 @@ class BrowserEnv(NetworkEnv):
     A NetworkEnv object that uses the BrowserSigner and BrowserRPC classes.
     """
 
-    def __init__(self, address=None, **kwargs):
-        super().__init__(rpc=BrowserRPC(), **kwargs)
-        self.signer = BrowserSigner(address)
+    def __init__(self, address=None, rpc=None, **kwargs):
+        if rpc is None:
+            rpc = BrowserRPC()
+        super().__init__(rpc=rpc, **kwargs)
+        self.signer = BrowserSigner(address, rpc)
         self.set_eoa(self.signer)
 
+    def get_chain_id(self) -> int:
+        chain_id: str = self._rpc.fetch("eth_chainId", [])
+        return int(chain_id, 0)
+
     def set_chain_id(self, chain_id: int | str):
         self._rpc.fetch(
             "wallet_switchEthereumChain",
@@ -162,7 +186,7 @@ def _javascript_call(js_func: str, *args, timeout_message: str) -> Any:
         logging.warning(f"Calling {js_func} with {args_str}")
 
     if colab_eval_js:
-        install_jupyter_javascript_triggers(BrowserRPC._debug_mode)
+        _install_javascript_triggers()
         result = colab_eval_js(js_code)
         return _parse_js_result(json.loads(result))
 
@@ -192,15 +216,15 @@ def _wait_buffer_set(buffer: memoryview, timeout_message: str) -> bytes:
     """
 
     async def _async_wait(deadline: float) -> bytes:
-        inner_loop = get_running_loop()
+        inner_loop = get_event_loop()
         while buffer.tobytes().startswith(NUL):
             if inner_loop.time() > deadline:
                 raise TimeoutError(timeout_message)
             await sleep(0.01)
 
         return buffer.tobytes().split(NUL)[0]
 
-    loop = get_running_loop()
+    loop = get_event_loop()
     future = _async_wait(deadline=loop.time() + CALLBACK_TOKEN_TIMEOUT.total_seconds())
     task = loop.create_task(future)
     loop.run_until_complete(task)

boa/integrations/jupyter/jupyter.js

@@ -5,9 +5,7 @@
 (() => {
     const rpc = async (method, params) => {
         const {ethereum} = window;
-        if (!ethereum) {
-            throw new Error('No Ethereum plugin found. Please authorize the site on your browser wallet.');
-        }
+        console.assert(ethereum, 'No Ethereum plugin found. Please authorize the site on your browser wallet.');
         return ethereum.request({method, params});
     };
 

tests/unitary/jupyter/test_browser.py

@@ -11,7 +11,7 @@
 from eth_account import Account
 
 import boa
-from boa.integrations.jupyter import BrowserSigner
+from boa.integrations.jupyter import BrowserRPC, BrowserSigner
 from boa.integrations.jupyter.browser import _generate_token
 from boa.rpc import RPCError
 
@@ -32,7 +32,7 @@ def mocked_token(token):
 
 @pytest.fixture()
 def env(account, mock_fork, mock_callback):
-    mock_callback("loadSigner", account.address)
+    mock_callback("eth_requestAccounts", [account.address])
     boa.set_browser_env(account)
     return boa.env
 
@@ -77,9 +77,7 @@ def find_response(mock_calls, func_to_body_dict):
 def mock_callback(mocked_token, display_mock):
     """Returns a function that allows mocking the result of the frontend callback."""
 
-    with mock.patch(
-        "boa.integrations.jupyter.browser.get_running_loop"
-    ) as mock_get_loop:
+    with mock.patch("boa.integrations.jupyter.browser.get_event_loop") as mock_get_loop:
         io_loop = mock_get_loop.return_value
         io_loop.time.return_value = 0
         func_to_body_dict = {}
@@ -153,27 +151,32 @@ def test_browser_sign_typed_data(display_mock, mock_callback, env):
     assert data == signature
 
 
+def test_browser_rpc_inject_js(mocked_token, display_mock, mock_callback):
+    _ = BrowserRPC()  # test the constructor
+    (((js1,), _),) = display_mock.call_args_list
+    assert "window._titanoboa = " in js1.data
+
+
 def test_browser_signer_colab(colab_eval_mock, mocked_token, display_mock):
     address = boa.env.generate_address()
-    colab_eval_mock.return_value = json.dumps({"data": address})
+    colab_eval_mock.return_value = json.dumps({"data": [address]})
     signer = BrowserSigner()
     assert signer.address == address
     colab_eval_mock.assert_called_once()
     (js,), _ = colab_eval_mock.call_args
-    assert f'loadSigner("{mocked_token}", null)' in js
+    assert f'rpc("{mocked_token}", "eth_requestAccounts", [])' in js
+    display_mock.assert_called_once()
 
 
 def test_browser_load_signer(token, display_mock, mock_callback, account, mock_fork):
-    mock_callback("loadSigner", account.address)
+    mock_callback("eth_requestAccounts", [account.address])
     boa.set_browser_env()
     assert boa.env.eoa == account.address
     assert type(boa.env._accounts[boa.env.eoa]).__name__ == BrowserSigner.__name__
 
 
 def test_browser_timeout():
-    with mock.patch(
-        "boa.integrations.jupyter.browser.get_running_loop"
-    ) as mock_get_loop:
+    with mock.patch("boa.integrations.jupyter.browser.get_event_loop") as mock_get_loop:
         io_loop = mock_get_loop.return_value
         now = datetime.now().timestamp()
         io_loop.time.side_effect = [now, now, now + 1000]
@@ -203,7 +206,8 @@ def dummy() -> bool:
     mock_callback("eth_chainId", "0x1")
     mock_callback("eth_estimateGas", "0x0")
     mock_callback("debug_traceCall", {})  # for prefetch state
-    mock_callback("sendTransaction", {"hash": "0x123"})
+    mock_callback("debug_traceCall", {})  # for prefetch state
+    mock_callback("eth_sendTransaction", {"hash": "0x123"})
     mock_callback(
         "waitForTransactionReceipt",
         {
@@ -231,11 +235,11 @@ def test_browser_chain_id(token, env, display_mock, mock_callback):
     )
 
 
-def test_browser_rpc(token, display_mock, mock_callback, account, mock_fork, env):
+def test_browser_rpc_gas_price(
+    token, display_mock, mock_callback, account, mock_fork, env
+):
     mock_callback("eth_gasPrice", "0x123")
     assert env.get_gas_price() == 291
-
-    assert display_mock.call_count == 7
     (js,), _ = display_mock.call_args
     assert f'rpc("{token}", "eth_gasPrice", [])' in js.data
     assert env._rpc.name == "BrowserRPC"
@@ -307,7 +311,9 @@ def test_browser_rpc_debug_error(mock_callback, env):
 
 
 def test_browser_js_error(token, display_mock, mock_callback, account, mock_fork):
-    mock_callback("loadSigner", error={"message": "custom message", "stack": ""})
+    mock_callback(
+        "eth_requestAccounts", error={"message": "custom message", "stack": ""}
+    )
     with pytest.raises(RPCError) as exc_info:
         BrowserSigner()
     assert str(exc_info.value) == "-1: custom message"