feat: default displayable trace in repr and str for CallTreeNode

.pre-commit-config.yaml

@@ -1,27 +1,27 @@
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.0.1
+    rev: v4.2.0
     hooks:
     -   id: check-yaml
 
 -   repo: https://github.com/pre-commit/mirrors-isort
-    rev: v5.9.3
+    rev: v5.10.1
     hooks:
       - id: isort
 
 -   repo: https://github.com/psf/black
-    rev: 21.10b0
+    rev: 22.3.0
     hooks:
       - id: black
         name: black
 
 -   repo: https://gitlab.com/pycqa/flake8
-    rev: 3.9.2
+    rev: 4.0.1
     hooks:
     -   id: flake8
 
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v0.910-1
+    rev: v0.960
     hooks:
     -   id: mypy
         additional_dependencies: [types-PyYAML, types-requests]

README.md

@@ -27,8 +27,52 @@ python3 setup.py install
 ```
 
 ## Quick Usage
-```bash
-ape console --network ethereum:local:hardhat
+
+If you are using a node that supports the `debug_traceTransaction` RPC, you can use `web3.py` to get trace frames:
+
+```python
+from web3 import HTTPProvider, Web3
+from evm_trace import TraceFrame
+
+web3 = Web3(HTTPProvider("https://path.to.my.node"))
+struct_logs = web3.manager.request_blocking("debug_traceTransaction", [txn_hash]).structLogs
+for item in struct_logs:
+    yield TraceFrame(**item)
+```
+
+If you want to get the call-tree node, you can do:
+
+```python
+from evm_trace import CallType, get_calltree_from_trace
+
+root_node_kwargs = {
+    "gas_cost": 10000000,
+    "gas_limit": 10000000000,
+    "address": "0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045",
+    "calldata": "0x00",
+    "value": 1000,
+    "call_type": CallType.MUTABLE,
+}
+
+# Where `trace` is a `TraceFrame` (see example above)
+calltree = get_calltree_from_trace(trace, **root_node_kwargs)
+```
+
+You can also customize the output by making your own display class:
+
+```python
+from evm_trace.display import DisplayableCallTreeNode, get_calltree_from_trace
+
+
+class CustomDisplay(DisplayableCallTreeNode):
+    def title(self) -> str:
+        call_type = self.call.call_type.value.lower().capitalize()
+        address = self.call.address.hex()
+        cost = self.call.gas_cost
+        return f"{call_type} call @ {address} gas_cost={cost}"
+
+
+calltree = get_calltree_from_trace(trace, display_cls=CustomDisplay)
 ```
 
 ## Development

evm_trace/base.py

@@ -1,11 +1,13 @@
 import math
-from enum import Enum
-from typing import Any, Dict, Iterator, List
+from typing import Any, Dict, Iterable, Iterator, List, Type
 
 from eth_utils import to_int
 from hexbytes import HexBytes
 from pydantic import BaseModel, Field, ValidationError, validator
 
+from evm_trace.display import DisplayableCallTreeNode
+from evm_trace.enums import CallType
+
 
 def _convert_hexbytes(cls, v: Any) -> HexBytes:
     try:
@@ -33,14 +35,6 @@ def convert_hexbytes_dict(cls, v) -> Dict[HexBytes, HexBytes]:
         return {_convert_hexbytes(cls, k): _convert_hexbytes(cls, val) for k, val in v.items()}
 
 
-class CallType(Enum):
-    INTERNAL = "INTERNAL"  # Non-opcode internal call
-    STATIC = "STATIC"  # STATICCALL opcode
-    MUTABLE = "MUTABLE"  # CALL opcode
-    DELEGATE = "DELEGATE"  # DELEGATECALL opcode
-    SELFDESTRUCT = "SELFDESTRUCT"  # SELFDESTRUCT opcode
-
-
 class CallTreeNode(BaseModel):
     call_type: CallType
     address: Any
@@ -53,11 +47,25 @@ class CallTreeNode(BaseModel):
     calls: List["CallTreeNode"] = []
     selfdestruct: bool = False
     failed: bool = False
+    display_cls: Type[DisplayableCallTreeNode] = DisplayableCallTreeNode
+
+    @property
+    def display_nodes(self) -> Iterable[DisplayableCallTreeNode]:
+        return self.display_cls.make_tree(self)
 
     @validator("address", "calldata", "returndata", pre=True)
     def validate_hexbytes(cls, v) -> HexBytes:
         return _convert_hexbytes(cls, v)
 
+    def __str__(self) -> str:
+        return "\n".join([str(t) for t in self.display_nodes])
+
+    def __repr__(self) -> str:
+        return str(self)
+
+    def __getitem__(self, index: int) -> "CallTreeNode":
+        return self.calls[index]
+
 
 def get_calltree_from_trace(
     trace: Iterator[TraceFrame], show_internal=False, **root_node_kwargs
@@ -68,17 +76,18 @@ def get_calltree_from_trace(
     Args:
         trace (Iterator[TraceFrame]): Iterator of transaction trace frames.
         show_internal (bool): Boolean whether to display internal calls. Defaulted to False.
-        root_node_kwargs (dict): Keyword argments passed to the root ``CallTreeNode``.
+        root_node_kwargs (dict): Keyword arguments passed to the root ``CallTreeNode``.
 
     Returns:
-        :class:`~evm_trace.base.CallTreeNode: Call tree of transaction trace.
+        :class:`~evm_trace.base.CallTreeNode`: Call tree of transaction trace.
     """
 
-    return _create_node_from_call(
+    node = _create_node_from_call(
         trace=trace,
         show_internal=show_internal,
         **root_node_kwargs,
     )
+    return node
 
 
 def _extract_memory(offset: HexBytes, size: HexBytes, memory: List[HexBytes]) -> HexBytes:
@@ -106,13 +115,15 @@ def _extract_memory(offset: HexBytes, size: HexBytes, memory: List[HexBytes]) ->
     # Compute the word that contains the last byte
     stop_word = math.ceil((offset_int + size_int) / 32)
 
-    byte_slice = b""
+    end_index = stop_word + 1
+    byte_slice = b"".join(memory[start_word:end_index])
+    offset_index = offset_int % 32
 
-    for word in memory[start_word:stop_word]:
-        byte_slice += word
+    # NOTE: Add 4 for the selector.
 
-    offset_index = offset_int % 32
-    return HexBytes(byte_slice[offset_index:size_int])
+    end_bytes_index = offset_index + size_int
+    return_bytes = byte_slice[offset_index:end_bytes_index]
+    return HexBytes(return_bytes)
 
 
 def _create_node_from_call(
@@ -127,16 +138,16 @@ def _create_node_from_call(
         raise NotImplementedError()
 
     node = CallTreeNode(**node_kwargs)
-
     for frame in trace:
         if frame.op in ("CALL", "DELEGATECALL", "STATICCALL"):
-            child_node_kwargs = {}
+            child_node_kwargs = {
+                "address": frame.stack[-2][-20:],
+                "depth": frame.depth,
+                "gas_limit": int(frame.stack[-1].hex(), 16),
+                "gas_cost": frame.gas_cost,
+            }
 
-            child_node_kwargs["address"] = frame.stack[-2][-20:]  # address is 20 bytes in EVM
-            child_node_kwargs["depth"] = frame.depth
             # TODO: Validate gas values
-            child_node_kwargs["gas_limit"] = int(frame.stack[-1].hex(), 16)
-            child_node_kwargs["gas_cost"] = frame.gas_cost
 
             if frame.op == "CALL":
                 child_node_kwargs["call_type"] = CallType.MUTABLE
@@ -160,10 +171,13 @@ def _create_node_from_call(
 
         # TODO: Handle internal nodes using JUMP and JUMPI
 
-        elif frame.op in ("SELFDESTRUCT", "STOP"):
-            # TODO: Handle the internal value transfer in SELFDESTRUCT
+        elif frame.op == "SELFDESTRUCT":
+            # TODO: Handle the internal value transfer
+            node.selfdestruct = True
+            break
+
+        elif frame.op == "STOP":
             # TODO: Handle "execution halted" vs. gas limit reached
-            node.selfdestruct = frame.op == "SELFDESTRUCT"
             break
 
         elif frame.op in ("RETURN", "REVERT"):
@@ -178,4 +192,5 @@ def _create_node_from_call(
         # NOTE: ignore other opcodes
 
     # TODO: Handle "execution halted" vs. gas limit reached
+
     return node

evm_trace/display.py

@@ -0,0 +1,83 @@
+from typing import TYPE_CHECKING, Iterable, Optional
+
+from eth_utils import to_checksum_address
+
+from evm_trace.enums import CallType
+
+if TYPE_CHECKING:
+    from evm_trace.base import CallTreeNode
+
+
+class DisplayableCallTreeNode(object):
+    _FILE_MIDDLE_PREFIX = "├──"
+    _FILE_LAST_PREFIX = "└──"
+    _PARENT_PREFIX_MIDDLE = "    "
+    _PARENT_PREFIX_LAST = "│   "
+
+    def __init__(
+        self,
+        call: "CallTreeNode",
+        parent: Optional["DisplayableCallTreeNode"] = None,
+        is_last: bool = False,
+    ):
+        self.call = call
+        self.parent = parent
+        self.is_last = is_last
+
+    @property
+    def depth(self) -> int:
+        return self.call.depth
+
+    @property
+    def title(self) -> str:
+        call_type = self.call.call_type.value.upper()
+        call_mnemonic = "CALL" if self.call.call_type == CallType.MUTABLE else f"{call_type}CALL"
+        address_hex_str = self.call.address.hex()
+
+        try:
+            address = to_checksum_address(address_hex_str)
+        except ImportError:
+            # Ignore checksumming if user does not have eth-hash backend installed.
+            address = address_hex_str
+
+        cost = self.call.gas_cost
+
+        call_path = str(address)
+        if self.call.calldata:
+            call_path = f"{call_path}.<{self.call.calldata[:4].hex()}>"
+
+        return f"{call_mnemonic}: {call_path} [{cost} gas]"
+
+    @classmethod
+    def make_tree(
+        cls,
+        root: "CallTreeNode",
+        parent: Optional["DisplayableCallTreeNode"] = None,
+        is_last: bool = False,
+    ) -> Iterable["DisplayableCallTreeNode"]:
+        displayable_root = cls(root, parent=parent, is_last=is_last)
+        yield displayable_root
+
+        count = 1
+        for child_node in root.calls:
+            is_last = count == len(root.calls)
+            if child_node.calls:
+                yield from cls.make_tree(child_node, parent=displayable_root, is_last=is_last)
+            else:
+                yield cls(child_node, parent=displayable_root, is_last=is_last)
+
+            count += 1
+
+    def __str__(self) -> str:
+        if self.parent is None:
+            return self.title
+
+        _filename_prefix = self._FILE_LAST_PREFIX if self.is_last else self._FILE_MIDDLE_PREFIX
+
+        parts = [f"{_filename_prefix} {self.title}"]
+        parent = self.parent
+        while parent and parent.parent is not None:
+            parts.append(self._PARENT_PREFIX_MIDDLE if parent.is_last else self._PARENT_PREFIX_LAST)
+            parent = parent.parent
+
+        return "".join(reversed(parts))

evm_trace/enums.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class CallType(Enum):
+    INTERNAL = "INTERNAL"  # Non-opcode internal call
+    STATIC = "STATIC"  # STATICCALL opcode
+    MUTABLE = "MUTABLE"  # CALL opcode
+    DELEGATE = "DELEGATE"  # DELEGATECALL opcode
+    SELFDESTRUCT = "SELFDESTRUCT"  # SELFDESTRUCT opcode

pyproject.toml

@@ -14,12 +14,11 @@ write_to = "evm_trace/version.py"
 
 [tool.black]
 line-length = 100
-target-version = ['py37', 'py38', 'py39']
+target-version = ['py37', 'py38', 'py39', 'py310']
 include = '\.pyi?$'
 
 [tool.pytest.ini_options]
 addopts = """
-    -n auto
     -p no:ape_test
     --cov-branch
     --cov-report term

setup.py

@@ -8,12 +8,13 @@
         "pytest-xdist",  # multi-process runner
         "pytest-cov",  # Coverage analyzer plugin
         "hypothesis>=6.2.0,<7.0",  # Strategy-based fuzzer
+        "eth-hash[pysha3]",  # For eth-utils address checksumming
     ],
     "lint": [
         "black>=22.3.0,<23.0",  # auto-formatter and linter
-        "mypy>=0.910,<1.0",  # Static type analyzer
-        "flake8>=3.8.3,<4.0",  # Style linter
-        "isort>=5.9.3,<6.0",  # Import sorting linter
+        "mypy>=0.960,<1.0",  # Static type analyzer
+        "flake8>=4.0.1,<5.0",  # Style linter
+        "isort>=5.10.1,<6.0",  # Import sorting linter
     ],
     "release": [  # `release` GitHub Action job uses this
         "setuptools",  # Installation tool
@@ -56,7 +57,7 @@
         "importlib-metadata ; python_version<'3.8'",
         "pydantic>=1.9.0,<2.0",
         "hexbytes>=0.2.2,<1.0.0",
-        "eth-utils==1.10.0",
+        "eth-utils>=1.10.0",
     ],  # NOTE: Add 3rd party libraries here
     python_requires=">=3.7.2,<4",
     extras_require=extras_require,
@@ -77,5 +78,6 @@
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
     ],
 )