feat: AST nodes [APE-967]

ape_solidity/_utils.py

@@ -20,7 +20,6 @@
 
 OUTPUT_SELECTION = [
     "abi",
-    "ast",
     "bin-runtime",
     "devdoc",
     "userdoc",

ape_solidity/compiler.py

@@ -11,12 +11,13 @@
 from ape.types import AddressType, ContractType
 from ape.utils import cached_property, get_relative_path
 from eth_utils import add_0x_prefix, is_0x_prefixed
-from ethpm_types import HexBytes, PackageManifest
+from ethpm_types import ASTNode, HexBytes, PackageManifest
+from ethpm_types.ast import ASTClassification
 from requests.exceptions import ConnectionError
 from semantic_version import NpmSpec, Version  # type: ignore
+from solcx import compile_standard  # type: ignore
 from solcx.exceptions import SolcError  # type: ignore
 from solcx.install import get_executable  # type: ignore
-from solcx.main import compile_standard  # type: ignore
 
 from ape_solidity._utils import (
     OUTPUT_SELECTION,
@@ -279,7 +280,8 @@ def get_compiler_settings(
             version_settings: Dict[str, Union[Any, List[Any]]] = {
                 "optimizer": {"enabled": self.config.optimize, "runs": 200},
                 "outputSelection": {
-                    str(get_relative_path(p, base_path)): {"*": OUTPUT_SELECTION} for p in sources
+                    str(get_relative_path(p, base_path)): {"*": OUTPUT_SELECTION, "": ["ast"]}
+                    for p in sources
                 },
             }
             remappings_used = set()
@@ -377,15 +379,26 @@ def compile(
                         if source_id in str(input_file_path):
                             input_contract_names.append(name)
 
+            def classify_ast(_node: ASTNode):
+                if _node.ast_type in ("FunctionDefinition", "FunctionDefinitionNode"):
+                    _node.classification = ASTClassification.FUNCTION
+
+                for child in _node.children:
+                    classify_ast(child)
+
             for source_id, contracts_out in output["contracts"].items():
-                for contract_name, contract_type in contracts_out.items():
+                ast_data = output["sources"][source_id]["ast"]
+                ast = ASTNode.parse_obj(ast_data)
+                classify_ast(ast)
+
+                for contract_name, ct_data in contracts_out.items():
                     contract_path = base_path / source_id
 
                     if contract_name not in input_contract_names:
                         # Only return ContractTypes explicitly asked for.
                         continue
 
-                    evm_data = contract_type["evm"]
+                    evm_data = ct_data["evm"]
 
                     # NOTE: This sounds backwards, but it isn't...
                     #  The "deployment_bytecode" is the same as the "bytecode",
@@ -413,17 +426,18 @@ def compile(
                                 ct for ct in contract_types if ct.name != contract_name
                             ]
 
-                    contract_type["contractName"] = contract_name
-                    contract_type["sourceId"] = str(
+                    ct_data["contractName"] = contract_name
+                    ct_data["sourceId"] = str(
                         get_relative_path(base_path / contract_path, base_path)
                     )
-                    contract_type["deploymentBytecode"] = {"bytecode": deployment_bytecode}
-                    contract_type["runtimeBytecode"] = {"bytecode": runtime_bytecode}
-                    contract_type["userdoc"] = load_dict(contract_type["userdoc"])
-                    contract_type["devdoc"] = load_dict(contract_type["devdoc"])
-                    contract_type["sourcemap"] = evm_data["bytecode"]["sourceMap"]
-                    contract_type_obj = ContractType.parse_obj(contract_type)
-                    contract_types.append(contract_type_obj)
+                    ct_data["deploymentBytecode"] = {"bytecode": deployment_bytecode}
+                    ct_data["runtimeBytecode"] = {"bytecode": runtime_bytecode}
+                    ct_data["userdoc"] = load_dict(ct_data["userdoc"])
+                    ct_data["devdoc"] = load_dict(ct_data["devdoc"])
+                    ct_data["sourcemap"] = evm_data["bytecode"]["sourceMap"]
+                    ct_data["ast"] = ast
+                    contract_type = ContractType.parse_obj(ct_data)
+                    contract_types.append(contract_type)
                     solc_versions_by_contract_name[contract_name] = solc_version
 
         return contract_types

setup.py

@@ -65,10 +65,11 @@
     include_package_data=True,
     install_requires=[
         "py-solc-x>=1.1.0,<2",
-        "eth-ape>=0.6.8,<0.7",
-        "ethpm-types",  # Use the version ape requires
+        "eth-ape>=0.6.9,<0.7",
+        "ethpm-types>=0.5.1,<0.6",  # Currently bumped 1 higher than eth-ape's.
         "packaging",  # Use the version ape requires
         "requests",
+        "typing-extensions==4.5.0",  # Can be removed onced pinned in Core.
     ],
     python_requires=">=3.8,<4",
     extras_require=extras_require,

tests/test_compiler.py

@@ -7,6 +7,7 @@
 import solcx  # type: ignore
 from ape.contracts import ContractContainer
 from ape.exceptions import CompilerError
+from ethpm_types.ast import ASTClassification
 from semantic_version import Version  # type: ignore
 
 from ape_solidity import Extension
@@ -330,8 +331,11 @@ def test_get_compiler_settings(compiler, project):
         output_selection = actual[version]["outputSelection"]
         assert actual[version]["optimizer"] == DEFAULT_OPTIMIZER
         for _, item_selection in output_selection.items():
-            for _, selection in item_selection.items():
-                assert selection == OUTPUT_SELECTION
+            for key, selection in item_selection.items():
+                if key == "*":  # All contracts
+                    assert selection == OUTPUT_SELECTION
+                elif key == "":  # All sources
+                    assert selection == ["ast"]
 
         actual_sources = [x for x in output_selection.keys()]
         for expected_source_id in expected_sources:
@@ -385,3 +389,11 @@ def test_enrich_error(compiler, project, owner, not_owner, connection):
         contract.withdraw(sender=not_owner)
 
     assert err.value.inputs == {"addr": not_owner.address, "counter": 123}
+
+
+def test_ast(project, compiler):
+    source_path = project.contracts_folder / "MultipleDefinitions.sol"
+    actual = compiler.compile([source_path])[-1].ast
+    fn_node = actual.children[1].children[0]
+    assert actual.ast_type == "SourceUnit"
+    assert fn_node.classification == ASTClassification.FUNCTION