Got overloaded function behavior working correctly

Also included a fix for https://github.com/bugout-dev/moonworm/issues/91
and https://github.com/bugout-dev/moonworm/issues/50

moonworm/abi.py

@@ -0,0 +1,78 @@
+"""
+ABI utilities, because web3 doesn't do selectors well.
+"""
+
+import glob
+import json
+import os
+from typing import Any, Dict, List, Optional
+
+from web3 import Web3
+
+
+def abi_input_signature(input_abi: Dict[str, Any]) -> str:
+    """
+    Stringifies a function ABI input object according to the ABI specification:
+    https://docs.soliditylang.org/en/v0.5.3/abi-spec.html
+    """
+    input_type = input_abi["type"]
+    if input_type.startswith("tuple"):
+        component_types = [
+            abi_input_signature(component) for component in input_abi["components"]
+        ]
+        input_type = f"({','.join(component_types)}){input_type[len('tuple'):]}"
+    return input_type
+
+
+def abi_function_signature(function_abi: Dict[str, Any]) -> str:
+    """
+    Stringifies a function ABI according to the ABI specification:
+    https://docs.soliditylang.org/en/v0.5.3/abi-spec.html
+    """
+    function_name = function_abi["name"]
+    function_arg_types = [
+        abi_input_signature(input_item) for input_item in function_abi["inputs"]
+    ]
+    function_signature = f"{function_name}({','.join(function_arg_types)})"
+    return function_signature
+
+
+def encode_function_signature(function_abi: Dict[str, Any]) -> Optional[str]:
+    """
+    Encodes the given function (from ABI) with arguments arg_1, ..., arg_n into its 4 byte signature
+    by calculating:
+    keccak256("<function_name>(<arg_1_type>,...,<arg_n_type>")
+
+    If function_abi is not actually a function ABI (detected by checking if function_abi["type"] == "function),
+    returns None.
+    """
+    if function_abi["type"] != "function":
+        return None
+    function_signature = abi_function_signature(function_abi)
+    encoded_signature = Web3.keccak(text=function_signature)[:4]
+    return encoded_signature.hex()
+
+
+def project_abis(project_dir: str) -> Dict[str, List[Dict[str, Any]]]:
+    """
+    Load all ABIs for project contracts and return then in a dictionary keyed by contract name.
+
+    Inputs:
+    - project_dir
+      Path to brownie project
+    """
+    build_dir = os.path.join(project_dir, "build", "contracts")
+    build_files = glob.glob(os.path.join(build_dir, "*.json"))
+
+    abis: Dict[str, List[Dict[str, Any]]] = {}
+
+    for filepath in build_files:
+        contract_name, _ = os.path.splitext(os.path.basename(filepath))
+        with open(filepath, "r") as ifp:
+            contract_artifact = json.load(ifp)
+
+        contract_abi = contract_artifact.get("abi", [])
+
+        abis[contract_name] = contract_abi
+
+    return abis

moonworm/generators/basic.py

@@ -17,6 +17,7 @@ import black.mode
 import inflection
 import libcst as cst
 
+from ..abi import encode_function_signature
 from ..version import MOONWORM_VERSION
 
 CONTRACT_TEMPLATE_PATH = os.path.join(os.path.dirname(__file__), "contract.py.template")
@@ -186,7 +187,9 @@ PROTECTED_ARG_NAMES: Set[str] = {
 }
 
 
-def function_spec(function_abi: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
+def function_spec(
+    function_abi: Dict[str, Any], is_overloaded: bool = False
+) -> Dict[str, Dict[str, Any]]:
     """
     Accepts function interface definitions from smart contract ABIs. An example input:
     {
@@ -226,6 +229,11 @@ def function_spec(function_abi: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
         ],
         "transact": False,
     }
+
+    If the caller provides `is_overloaded`, the function signature is suffixed to the method and CLI
+    names with the appropriate dashes. No changes are made to the ABI name itself. For example, generated
+    brownie interfaces rely on brownie's internal deduplication logic (based on argument types) to delegate
+    behavior to the right version of the brownie contract method.
     """
     abi_name = function_abi.get("name")
     if abi_name is None:
@@ -235,6 +243,11 @@ def function_spec(function_abi: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
     function_name = normalize_abi_name(underscored_name)
     cli_name = inflection.dasherize(underscored_name)
 
+    if is_overloaded:
+        signature = encode_function_signature(function_abi)
+        function_name += f"_{signature}"
+        cli_name += f"-{signature}"
+
     default_input_name = "arg"
     default_counter = 1
 
@@ -275,7 +288,8 @@ def function_spec(function_abi: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
         inputs.append(input_spec)
 
     transact = True
-    if function_abi.get("stateMutability") == "view":
+    state_mutability = function_abi.get("stateMutability", "").lower()
+    if state_mutability == "view" or state_mutability == "pure":
         transact = False
 
     spec = {

moonworm/generators/brownie.py

@@ -7,7 +7,7 @@ The entrypoint to code generation is [`generate_brownie_interface`][moonworm.gen
 import copy
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Set
 
 import libcst as cst
 from libcst._nodes.statement import SimpleStatementLine
@@ -33,6 +33,21 @@ except Exception as e:
     logging.warn(e)
 
 
+def get_overloaded_functions(abi: List[Dict[str, Any]]) -> Set[str]:
+    """
+    Return a set containing the function names of all overloaded functions.
+    """
+    function_name_counters: Dict[str, int] = {}
+    for item in abi:
+        if item["type"] != "function":
+            continue
+        if item["name"] in function_name_counters:
+            function_name_counters[item["name"]] += 1
+        else:
+            function_name_counters[item["name"]] = 1
+    return {name for name, count in function_name_counters.items() if count > 1}
+
+
 def generate_brownie_contract_class(
     abi: List[Dict[str, Any]],
     contract_name: str,
@@ -78,6 +93,8 @@ def generate_brownie_contract_class(
     contract_constructor = get_constructor(abi)
     contract_constructor["name"] = "constructor"
 
+    overloaded_functions = get_overloaded_functions(abi)
+
     class_functions = (
         [class_constructor]
         + [
@@ -86,7 +103,9 @@ def generate_brownie_contract_class(
             generate_verify_contract(),
         ]
         + [
-            generate_brownie_contract_function(function)
+            generate_brownie_contract_function(
+                function, function.get("name", "") in overloaded_functions
+            )
             for function in abi
             if function["type"] == "function"
         ]
@@ -187,8 +206,10 @@ def generate_assert_contract_is_instantiated() -> cst.FunctionDef:
     return function_def
 
 
-def generate_brownie_contract_function(func_object: Dict[str, Any]) -> cst.FunctionDef:
-    spec = function_spec(func_object)
+def generate_brownie_contract_function(
+    func_object: Dict[str, Any], is_overloaded: bool = False
+) -> cst.FunctionDef:
+    spec = function_spec(func_object, is_overloaded)
     func_params = []
     func_params.append(cst.Param(name=cst.Name("self")))
 
@@ -516,7 +537,7 @@ def generate_verify_contract_handler(contract_name: str) -> Optional[cst.Functio
 
 
 def generate_cli_handler(
-    function_abi: Dict[str, Any], contract_name: str
+    function_abi: Dict[str, Any], contract_name: str, is_overloaded: bool = False
 ) -> Optional[cst.FunctionDef]:
     """
     Generates a handler which translates parsed command line arguments to method calls on the generated
@@ -525,7 +546,7 @@ def generate_cli_handler(
     Returns None if it is not appropriate for the given function to have a handler (e.g. fallback or
     receive). constructor is handled separately with a deploy handler.
     """
-    spec = function_spec(function_abi)
+    spec = function_spec(function_abi, is_overloaded)
     function_name = spec["method"]
 
     function_body_raw: List[cst.CSTNode] = []
@@ -734,8 +755,16 @@ def generate_cli_generator(
         "inputs": [],
     }
 
+    overloaded_functions = get_overloaded_functions(abi)
+
     specs: List[Dict[str, Any]] = [constructor_spec, verify_contract_spec]
-    specs.extend([function_spec(item) for item in abi if item["type"] == "function"])
+    specs.extend(
+        [
+            function_spec(item, item.get("name", "") in overloaded_functions)
+            for item in abi
+            if item["type"] == "function"
+        ]
+    )
 
     for spec in specs:
         subparser_statements: List[SimpleStatementLine] = [cst.Newline()]
@@ -893,9 +922,11 @@ def generate_brownie_cli(
         add_deploy_handler,
         add_verify_contract_handler,
     ]
+
+    overloaded_functions = get_overloaded_functions(abi)
     handlers.extend(
         [
-            generate_cli_handler(function_abi, contract_name)
+            generate_cli_handler(function_abi, contract_name, function_abi.get("name") in overloaded_functions)
             for function_abi in abi
             if function_abi.get("type") == "function"
             and function_abi.get("name") is not None

moonworm/version.py

@@ -1 +1 @@
-MOONWORM_VERSION = "0.5.3"
+MOONWORM_VERSION = "0.6.0"