Merge pull request #163 from DanielSchiavini/perf

fix: Improve Jupyter performance

boa/integrations/jupyter/browser.py

@@ -29,6 +29,7 @@
 )
 from boa.network import NetworkEnv
 from boa.rpc import RPC, RPCError
+from boa.util.abi import Address
 
 try:
     from google.colab.output import eval_js as colab_eval_js
@@ -39,6 +40,12 @@
 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()
+
+
 class BrowserSigner:
     """
     A BrowserSigner is a class that can be used to sign transactions in IPython/JupyterLab.
@@ -49,12 +56,10 @@ def __init__(self, address=None):
         Create a BrowserSigner instance.
         :param address: The account address. If not provided, it will be requested from the browser.
         """
-        if address is not None:
-            self.address = address
-        else:
-            self.address = _javascript_call(
-                "loadSigner", timeout_message=ADDRESS_TIMEOUT_MESSAGE
-            )
+        address = _javascript_call(
+            "loadSigner", address, timeout_message=ADDRESS_TIMEOUT_MESSAGE
+        )
+        self.address = Address(address)
 
     def send_transaction(self, tx_data: dict) -> dict:
         """
@@ -161,22 +166,22 @@ def _javascript_call(js_func: str, *args, timeout_message: str) -> Any:
     :param kwargs: The arguments to pass to the Javascript snippet.
     :return: The result of the Javascript snippet sent to the API.
     """
-    install_jupyter_javascript_triggers()
-
     token = _generate_token()
     args_str = ", ".join(json.dumps(p) for p in chain([token], args))
-    js_code = f"window._titanoboa.{js_func}({args_str})"
+    js_code = f"window._titanoboa.{js_func}({args_str});"
     # logging.warning(f"Calling {js_func} with {args_str}")
 
     if colab_eval_js:
+        install_jupyter_javascript_triggers()
         result = colab_eval_js(js_code)
         return _parse_js_result(json.loads(result))
 
     memory = SharedMemory(name=token, create=True, size=SHARED_MEMORY_LENGTH)
     logging.info(f"Waiting for {token}")
     try:
         memory.buf[:1] = NUL
-        display(Javascript(js_code))
+        hide_output_element = "element.style.display = 'none';"
+        display(Javascript(js_code + hide_output_element))
         message_bytes = _wait_buffer_set(memory.buf, timeout_message)
         return _parse_js_result(json.loads(message_bytes.decode()))
     finally:

boa/integrations/jupyter/constants.py

@@ -4,7 +4,6 @@
 CALLBACK_TOKEN_TIMEOUT = timedelta(minutes=3)
 SHARED_MEMORY_LENGTH = 50 * 1024 + len(NUL)  # Size of the shared memory object
 CALLBACK_TOKEN_BYTES = 32
-ETHERS_JS_URL = "https://cdnjs.cloudflare.com/ajax/libs/ethers/6.9.0/ethers.umd.min.js"
 PLUGIN_NAME = "titanoboa_jupyterlab"
 TOKEN_REGEX = rf"{PLUGIN_NAME}_[0-9a-fA-F]{{{CALLBACK_TOKEN_BYTES * 2}}}"
 TRANSACTION_TIMEOUT_MESSAGE = (

boa/integrations/jupyter/jupyter.js

@@ -3,14 +3,12 @@
  * BrowserSigner to the frontend.
  */
 (() => {
-    let provider; // cache the provider to avoid re-creating it every time
-    const getEthersProvider = () => {
-        if (provider) return provider;
+    const rpc = (method, params) => {
         const {ethereum} = window;
         if (!ethereum) {
             throw new Error('No Ethereum plugin found. Please authorize the site on your browser wallet.');
         }
-        return provider = new ethers.BrowserProvider(ethereum);
+        return ethereum.request({method, params});
     };
 
     /** Stringify data, converting big ints to strings */
@@ -40,19 +38,21 @@
         return response.text();
     }
 
-    const getSigner = () => getEthersProvider().getSigner();
-
-    /** Load the signer via ethers user */
-    const loadSigner = () => getSigner().then(s => s.getAddress());
+    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 = transaction => getSigner().then(s => s.sendTransaction(transaction));
+    const sendTransaction = async transaction => ({"hash": await rpc('eth_sendTransaction', [transaction])});
 
     /** Sign a typed data via ethers */
-    const signTypedData = (domain, types, value) => getSigner().then(s => s.signTypedData(domain, types, value));
-
-    /** Call an RPC method via ethers */
-    const rpc = (method, params) => getEthersProvider().send(method, params);
+    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) => {
@@ -107,4 +107,6 @@
         rpc: handleCallback(rpc),
         multiRpc: handleCallback(multiRpc),
     };
+
+    if (element) element.style.display = "none";  // hide the output element in JupyterLab
 })();

boa/integrations/jupyter/utils.py

@@ -1,20 +1,14 @@
 from os.path import dirname, join, realpath
 
-import requests
 from IPython.display import Javascript, display
 
-from boa.integrations.jupyter.constants import ETHERS_JS_URL
-
 
 def install_jupyter_javascript_triggers():
     """Run the ethers and titanoboa_jupyterlab Javascript snippets in the browser."""
-    ethers_js = requests.get(ETHERS_JS_URL)
-
     cur_dir = dirname(realpath(__file__))
     with open(join(cur_dir, "jupyter.js")) as f:
         jupyter_js = f.read()
-
-    display(Javascript(ethers_js.text + jupyter_js))
+    display(Javascript(jupyter_js))
 
 
 def convert_frontend_dict(data):