Revert "Revert code changes"

This reverts commit 50ee8e9

boa/integrations/jupyter/browser.py

@@ -4,10 +4,11 @@
 """
 import json
 import logging
-from asyncio import get_running_loop, sleep
+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
@@ -23,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
@@ -40,26 +37,74 @@
 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 js_file:
+        display(Javascript(js_file.read()))
+
+
+class BrowserRPC(RPC):
+    """
+    An RPC object that sends requests to the browser via Javascript.
+    """
+
+    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=BrowserRPC()):
         """
         Create a BrowserSigner instance.
         :param address: The account address. If not provided, it will be requested from the browser.
         """
+        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 accounts:
+            address = accounts[0]
+        elif 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:
@@ -70,10 +115,10 @@ 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, domain: dict[str, Any], types: dict[str, list], value: dict[str, Any]
@@ -85,48 +130,9 @@ def sign_typed_data(
         :param value: The value to sign.
         :return: The signature.
         """
-        return _javascript_call(
-            "signTypedData",
-            domain,
-            types,
-            value,
-            timeout_message=TRANSACTION_TIMEOUT_MESSAGE,
-        )
-
-
-class BrowserRPC(RPC):
-    """
-    An RPC object that sends requests to the browser via Javascript.
-    """
-
-    @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,
+        payload = json.dumps({"domain": domain, "types": types, "value": value})
+        return self._rpc.fetch(
+            "eth_signTypedData_v4", [self.address, payload], TRANSACTION_TIMEOUT_MESSAGE
         )
 
 
@@ -135,23 +141,19 @@ 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)
+    def __init__(self, address=None, rpc=BrowserRPC()):
+        super().__init__(rpc=rpc)
+        self.signer = BrowserSigner(address, rpc)
         self.set_eoa(self.signer)
 
     def get_chain_id(self) -> int:
-        chain_id = _javascript_call(
-            "rpc", "eth_chainId", timeout_message=RPC_TIMEOUT_MESSAGE
-        )
-        return int.from_bytes(bytes.fromhex(chain_id[2:]), "big")
+        chain_id: str = self._rpc.fetch("eth_chainId", [])
+        return int(chain_id, 0)
 
     def set_chain_id(self, chain_id: int | str):
-        _javascript_call(
-            "rpc",
+        self._rpc.fetch(
             "wallet_switchEthereumChain",
             [{"chainId": chain_id if isinstance(chain_id, str) else hex(chain_id)}],
-            timeout_message=RPC_TIMEOUT_MESSAGE,
         )
         self._reset_fork()
 
@@ -173,7 +175,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()
+        _install_javascript_triggers()
         result = colab_eval_js(js_code)
         return _parse_js_result(json.loads(result))
 
@@ -203,15 +205,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 = (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});
     };
 
@@ -38,22 +36,6 @@
         return response.text();
     }
 
-    let from;
-    const loadSigner = async (address) => {
-        const accounts = await rpc('eth_requestAccounts');
-        from = accounts.includes(address) ? address : accounts[0];
-        return from;
-    };
-
-    /** Sign a transaction via ethers */
-    const sendTransaction = async transaction => ({"hash": await rpc('eth_sendTransaction', [transaction])});
-
-    /** Sign a typed data via ethers */
-    const signTypedData = (domain, types, value) => rpc(
-        'eth_signTypedData_v4',
-        [from, JSON.stringify({domain, types, value})]
-    );
-
     /** Wait until the transaction is mined */
     const waitForTransactionReceipt = async (tx_hash, timeout, poll_latency) => {
         while (true) {
@@ -100,9 +82,6 @@
 
     // expose functions to window, so they can be called from the BrowserSigner
     window._titanoboa = {
-        loadSigner: handleCallback(loadSigner),
-        sendTransaction: handleCallback(sendTransaction),
-        signTypedData: handleCallback(signTypedData),
         waitForTransactionReceipt: handleCallback(waitForTransactionReceipt),
         rpc: handleCallback(rpc),
         multiRpc: handleCallback(multiRpc),

tests/unitary/jupyter/test_browser.py

@@ -10,7 +10,7 @@
 from eth_account import Account
 
 import boa
-from boa.integrations.jupyter import BrowserSigner
+from boa.integrations.jupyter import BrowserRPC, BrowserSigner
 from boa.rpc import RPCError
 
 
@@ -25,7 +25,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
 
@@ -70,9 +70,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 = {}
@@ -135,25 +133,32 @@ def test_nest_applied():
 
 def test_browser_sign_typed_data(display_mock, mock_callback, env):
     signature = env.generate_address()
-    mock_callback("signTypedData", signature)
+    mock_callback("eth_signTypedData_v4", signature)
     data = env.signer.sign_typed_data(
         {"name": "My App"}, {"types": []}, {"data": "0x1234"}
     )
     assert data == signature
 
 
+def test_browser_rpc_inject_js(mocked_token, display_mock, mock_callback):
+    BrowserRPC()
+    (((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_loads_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__
@@ -206,7 +211,9 @@ def test_browser_rpc_server_error(
 
 
 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) == "CALLBACK_ERROR: custom message"