Expose install, allow getting whole arrays

boa_zksync/browser.py

@@ -7,6 +7,7 @@
 from boa.rpc import EthereumRPC
 
 from boa_zksync.environment import ZksyncEnv
+from boa_zksync.util import install_zkvyper_compiler, install_era_test_node
 
 
 class ZksyncBrowserEnv(ZksyncEnv):
@@ -42,45 +43,3 @@ def fork_rpc(
             install_era_test_node()
 
         return super().fork_rpc(rpc, reset_traces, block_identifier, **kwargs)
-
-
-def install_zkvyper_compiler(
-    source="https://raw.githubusercontent.com/matter-labs/zkvyper-bin/"
-    "66cc159d9b6af3b5616f6ed7199bd817bf42bf0a/linux-amd64/zkvyper-linux-amd64-musl-v1.4.0",
-    destination="/usr/local/bin/zkvyper",
-):
-    """
-    Downloads the zkvyper binary from the given source URL and installs it to
-    the destination directory.
-
-    This is a very basic implementation - usually users want to install the binary
-    manually, but in the Colab environment, we can automate this process.
-    """
-    response = requests.get(source)
-    with open(destination, "wb") as f:
-        f.write(response.content)
-
-    os.chmod(destination, 0o755)  # make it executable
-    assert os.system("zkvyper --version") == 0  # check if it works
-
-
-def install_era_test_node(
-    source="https://github.com/matter-labs/era-test-node/releases/download/"
-    "v0.1.0-alpha.19/era_test_node-v0.1.0-alpha.19-x86_64-unknown-linux-gnu.tar.gz",
-    destination="/usr/local/bin/era_test_node",
-):
-    """
-    Downloads the era-test-node binary from the given source URL and installs it to
-    the destination directory.
-
-    This is a very basic implementation - usually users want to install the binary
-    manually, but in the Colab environment, we can automate this process.
-    """
-    response = requests.get(source)
-    with open("era_test_node.tar.gz", "wb") as f:
-        f.write(response.content)
-
-    os.system("tar --extract --file=era_test_node.tar.gz")
-    os.system(f"mv era_test_node {destination}")
-    os.system(f"{destination} --version")
-    os.system("rm era_test_node.tar.gz")

boa_zksync/contract.py

@@ -12,6 +12,7 @@
 from boa.util.abi import Address
 from cached_property import cached_property
 from vyper.semantics.analysis.base import VarInfo
+from vyper.semantics.types import HashMapT
 from vyper.semantics.types.function import ContractFunctionT
 
 from boa_zksync.compile import compile_zksync_source
@@ -143,7 +144,10 @@ def source_code(self):
 
 class ZksyncInternalVariable(_ZksyncInternal):
     def __init__(self, var: VarInfo, name: str, contract: ZksyncContract):
-        inputs, output = var.typ.getter_signature
+        if isinstance(var.typ, HashMapT):
+            inputs, output = var.typ.getter_signature
+        else:
+            inputs, output = [], var.typ
         abi = {
             "anonymous": False,
             "inputs": [
@@ -165,18 +169,15 @@ def get(self, *args):
 
     @cached_property
     def source_code(self):
-        args, arg_getter = "", ""
-        inputs, output = self.var.typ.getter_signature
-        if inputs:
-            arg_getter = "".join([f"[arg{i}]" for i in range(len(inputs))])
-            args = ", ".join([f"arg{i}: {arg.abi_type.selector_name()}" for i, arg in enumerate(inputs)])
-
+        inputs, output_type = self._abi["inputs"], self.return_type[0]
+        getter_call = "".join(f"[{i['name']}]" for i in inputs)
+        args_signature = ", ".join(f"{i['name']}: {i['type']}" for i in inputs)
         return textwrap.dedent(
             f"""
             @external
             @payable
-            def __boa_private_{self.var_name}__({args}) -> {output.abi_type.selector_name()}:
-                return self.{self.var_name}{arg_getter}
+            def __boa_private_{self.var_name}__({args_signature}) -> {output_type}:
+                return self.{self.var_name}{getter_call}
         """
         )
 

boa_zksync/util.py

@@ -1,3 +1,4 @@
+import os
 import socket
 from datetime import datetime, timedelta
 from subprocess import Popen, TimeoutExpired
@@ -35,3 +36,45 @@ def stop_subprocess(proc: Popen[bytes]):
     except TimeoutExpired:  # pragma: no cover
         proc.kill()
         proc.wait(timeout=1)
+
+
+def install_zkvyper_compiler(
+    source="https://raw.githubusercontent.com/matter-labs/zkvyper-bin/"
+    "66cc159d9b6af3b5616f6ed7199bd817bf42bf0a/linux-amd64/zkvyper-linux-amd64-musl-v1.4.0",
+    destination="/usr/local/bin/zkvyper",
+):
+    """
+    Downloads the zkvyper binary from the given source URL and installs it to
+    the destination directory.
+
+    This is a very basic implementation - usually users want to install the binary
+    manually, but in the Colab environment, we can automate this process.
+    """
+    response = requests.get(source)
+    with open(destination, "wb") as f:
+        f.write(response.content)
+
+    os.chmod(destination, 0o755)  # make it executable
+    assert os.system("zkvyper --version") == 0  # check if it works
+
+
+def install_era_test_node(
+    source="https://github.com/matter-labs/era-test-node/releases/download/"
+    "v0.1.0-alpha.19/era_test_node-v0.1.0-alpha.19-x86_64-unknown-linux-gnu.tar.gz",
+    destination="/usr/local/bin/era_test_node",
+):
+    """
+    Downloads the era-test-node binary from the given source URL and installs it to
+    the destination directory.
+
+    This is a very basic implementation - usually users want to install the binary
+    manually, but in the Colab environment, we can automate this process.
+    """
+    response = requests.get(source)
+    with open("era_test_node.tar.gz", "wb") as f:
+        f.write(response.content)
+
+    os.system("tar --extract --file=era_test_node.tar.gz")
+    os.system(f"mv era_test_node {destination}")
+    os.system(f"{destination} --version")
+    os.system("rm era_test_node.tar.gz")

tests/test_deploy.py

@@ -170,17 +170,23 @@ def test_private(zksync_env):
     code = """
 bar: uint256
 map: HashMap[uint256, uint256]
+list: uint256[2]
 
 @internal
 def foo(x: uint256) -> uint256:
     self.bar = x
+    self.map[0] = x
+    self.list[0] = x
     return x
 """
     contract = boa.loads(code)
     assert contract._storage.bar.get() == 0
     assert contract._storage.map.get(0) == 0
+    assert contract._storage.list.get() == [0, 0]
     assert contract.internal.foo(123) == 123
     assert contract._storage.bar.get() == 123
+    assert contract._storage.map.get(0) == 123
+    assert contract._storage.list.get() == [123, 0]
     assert contract.eval("self.bar = 456") is None
     assert contract.eval("self.bar") == 456