[wip] add Vm.vy

foundry.toml

@@ -1,5 +1,6 @@
 [profile.default]
 fs_permissions = [{ access = "read-write", path = "./"}]
+skip = ["src/Vm.vy"]
 
 [rpc_endpoints]
 # The RPC URLs are modified versions of the default for testing initialization.

scripts/vm_vyper.py

@@ -0,0 +1,202 @@
+#!/usr/bin/env python3
+
+from vm import *
+
+OUT_PATH = "src/Vm.vy"
+
+def main():
+    json_str = request.urlopen(CHEATCODES_JSON_URL).read().decode("utf-8")
+    contract = Cheatcodes.from_json(json_str)
+
+    ccs = contract.cheatcodes
+    ccs = list(filter(lambda cc: cc.status not in ["experimental", "internal"], ccs))
+    ccs.sort(key=lambda cc: cc.func.id)
+
+    out = ""
+
+    out += "# Automatically @generated by scripts/vm_vyper.py. Do not modify manually.\n\n"
+
+    pp = VyperCheatcodesPrinter(
+        spdx_identifier="MIT OR Apache-2.0",
+        solidity_requirement=">=0.6.2 <0.9.0",
+        abicoder_pragma=True,
+    )
+    pp.prelude = False
+    out += pp.finish()
+
+    vm = Cheatcodes(
+        # TODO: Custom errors were introduced in 0.8.4
+        errors=[],  # contract.errors
+        events=contract.events,
+        enums=contract.enums,
+        structs=contract.structs,
+        cheatcodes=filter_functions_for_vyper(ccs),
+    )
+    pp.p_contract(vm, "Vm")
+    out += pp.finish()
+
+    out += "\n\nvm: constant(Vm) = Vm(0x7109709ECfa91a80626fF3989D68f67F5b1DD12D)\n"
+
+    with open(OUT_PATH, "w") as f:
+        f.write(out)
+
+    print(f"Wrote to {OUT_PATH}")
+
+# Converts Solidity type name to Vyper type name
+def to_vyper_ty(ty: str) -> str:
+    if ty.endswith("[]"):
+        return f"DynArray[{to_vyper_ty(ty[:-2])}, 1024]"
+
+    ty = ty.replace("string", "String[1024]")
+    if ty == "bytes":
+        ty = "Bytes[1024]"
+    return ty
+
+# Fixes identifiers to not conflict with Vyper keywords
+def fix_vyper_identifier(s: str) -> str:
+    if s in ["from", "None", "value", "gas"]:
+        return s + "_"
+    else:
+        return s
+
+# Converts Solidity function declaration to Vyper function declaration
+def to_vyper_func(declaration: str) -> str:
+    func_name = declaration.split("function")[1].split("(")[0].replace(" ", "")
+    params_list = declaration.split("(")[1].split(")")[0]
+
+    if "returns" in declaration:
+        ret_params_list = declaration.split("returns")[1].split("(")[1].split(")")[0]
+    else:
+        ret_params_list = ""
+
+    if params_list != "":
+        params_list = [param.replace("memory", "").replace("calldata", "") for param in params_list.split(",")]
+        params_list = ", ".join([f"{fix_vyper_identifier(param.split()[1])}: {to_vyper_ty(param.split()[0])}" for param in params_list])
+
+    if ret_params_list != "":
+        ret_params = [to_vyper_ty(param.split()[0]) for param in ret_params_list.split(",")]
+        if len(ret_params) == 1:
+            ret_params_list = ret_params[0]
+        else:
+            ret_params_list = f'({", ".join(ret_params)})'
+
+    if "view" in declaration:
+        visibility = "view"
+    else:
+        visibility = "nonpayable"
+
+    declaration = f"def {func_name}({params_list})"
+
+    if ret_params_list != "":
+        declaration += f" -> {ret_params_list}"
+
+    declaration += f": {visibility}"
+
+    return declaration
+
+# Filters function list before converting to Vyper code.
+# Excludes conflicting overloads
+def filter_functions_for_vyper(functions: list['Cheatcode']) -> list[str]:
+    functions = list(filter(lambda f: f.func.id != "rpcUrls", functions))
+    name_to_funcs = {}
+    for func in functions:
+        name = func.func.declaration.split("function")[1].split("(")[0].strip()
+        name_to_funcs.setdefault(name, []).append(func)
+
+    new_funcs = []
+    for name, funcs in name_to_funcs.items():
+        new_funcs.append(funcs[0])
+
+    return new_funcs
+
+class VyperCheatcodesPrinter(CheatcodesPrinter):
+    def p_contract(self, contract: Cheatcodes, name: str, inherits: str = ""):
+        if self.prelude:
+            self.p_prelude(contract)
+
+        for item in self.items_order.get_list():
+            if item == Item.EVENT:
+                continue
+            elif item == Item.ERROR:
+                continue
+            elif item == Item.ENUM:
+                self.p_enums(contract.enums)
+            elif item == Item.STRUCT:
+                self.p_structs(contract.structs)
+            elif item == Item.FUNCTION:
+                self._p_str("interface ")
+                name = name.strip()
+                if name != "":
+                    self._p_str(name)
+                self._p_str(":")
+                self._p_nl()
+                self._with_indent(lambda: self.p_functions(contract.cheatcodes))
+                self._p_nl()
+            else:
+                assert False, f"unknown item {item}"
+
+    def p_struct(self, struct: Struct):
+        self._p_comment(struct.description, doc=True)
+        self._p_line(lambda: self._p_str(f"struct {struct.name}:"))
+        self._with_indent(lambda: self.p_struct_fields(struct.fields))
+
+    def p_struct_field(self, field: StructField):
+        self._p_comment(field.description)
+        self._p_indented(lambda: self._p_str(f"{fix_vyper_identifier(field.name)}: {to_vyper_ty(field.ty)}"))
+
+    def p_enum(self, enum: Enum):
+        self._p_comment(enum.description, doc=True)
+        self._p_line(lambda: self._p_str(f"enum {enum.name}:"))
+        self._with_indent(lambda: self.p_enum_variants(enum.variants))
+
+    def p_enum_variants(self, variants: list[EnumVariant]):
+        for variant in variants:
+            self._p_indent()
+            self._p_comment(variant.description)
+
+            self._p_indent()
+            self._p_str(fix_vyper_identifier(variant.name))
+            self._p_nl()
+
+    def p_function(self, func: Function):
+        self._p_comment(func.description, doc=True)
+        self._p_line(lambda: self._p_str(to_vyper_func(func.declaration)))
+
+    def _p_comment(self, s: str, doc: bool = False):
+        s = s.strip()
+        if s == "":
+            return
+
+        s = map(lambda line: line.lstrip(), s.split("\n"))
+        if self.block_doc_style:
+            self._p_str("/*")
+            if doc:
+                self._p_str("*")
+            self._p_nl()
+            for line in s:
+                self._p_indent()
+                self._p_str(" ")
+                if doc:
+                    self._p_str("* ")
+                self._p_str(line)
+                self._p_nl()
+            self._p_indent()
+            self._p_str(" */")
+            self._p_nl()
+        else:
+            first_line = True
+            for line in s:
+                if not first_line:
+                    self._p_indent()
+                first_line = False
+
+                if doc:
+                    self._p_str("# ")
+                else:
+                    self._p_str("# ")
+                self._p_str(line)
+                self._p_nl()
+
+
+if __name__ == "__main__":
+    main()