Change NestedTrace to use ClassifiedTraces

mev_inspect/schemas/__init__.py

@@ -1,2 +1,2 @@
 from .abi import ABI
-from .blocks import Block, NestedTrace, Trace, TraceType
+from .blocks import Block, Trace, TraceType

mev_inspect/schemas/blocks.py

@@ -68,9 +68,3 @@ class Block(Web3Model):
         return [trace for trace in self.traces if trace.transaction_hash == hash]
 
 
-class NestedTrace(BaseModel):
-    trace: Trace
-    subtraces: List["NestedTrace"]
-
-
-NestedTrace.update_forward_refs()

mev_inspect/schemas/classified_traces.py

@@ -50,3 +50,11 @@ class ClassifierSpec(BaseModel):
     protocol: Optional[Protocol] = None
     valid_contract_addresses: Optional[List[str]] = None
     classifications: Dict[str, Classification] = {}
+
+
+class NestedTrace(BaseModel):
+    trace: ClassifiedTrace
+    subtraces: List["NestedTrace"]
+
+
+NestedTrace.update_forward_refs()

mev_inspect/traces.py

@@ -1,10 +1,10 @@
 from itertools import groupby
 from typing import Iterable, List
 
-from mev_inspect.schemas import Trace, NestedTrace
+from mev_inspect.schemas.classified_traces import ClassifiedTrace, NestedTrace
 
 
-def as_nested_traces(traces: Iterable[Trace]) -> List[NestedTrace]:
+def as_nested_traces(traces: Iterable[ClassifiedTrace]) -> List[NestedTrace]:
     nested_traces = []
 
     sorted_by_transaction_hash = sorted(traces, key=_get_transaction_hash)
@@ -20,7 +20,7 @@ def _get_transaction_hash(trace) -> str:
     return trace.transaction_hash
 
 
-def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrace]:
+def _as_nested_traces_by_transaction(traces: Iterable[ClassifiedTrace]) -> List[NestedTrace]:
     """
     Turns a list of Traces into a a tree of NestedTraces
     using their trace addresses
@@ -35,7 +35,7 @@ def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrac
     nested_traces = []
 
     parent = None
-    children: List[Trace] = []
+    children: List[ClassifiedTrace] = []
 
     sorted_traces = sorted(traces, key=lambda t: t.trace_address)
 
@@ -70,7 +70,7 @@ def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrac
     return nested_traces
 
 
-def _is_subtrace(trace: Trace, parent: Trace):
+def _is_subtrace(trace: ClassifiedTrace, parent: ClassifiedTrace):
     parent_trace_length = len(parent.trace_address)
 
     if len(trace.trace_address) > parent_trace_length:

tests/trace_test.py

@@ -1,7 +1,8 @@
 import unittest
 from typing import List
 
-from mev_inspect.schemas import Trace, TraceType, NestedTrace
+from mev_inspect.schemas import TraceType
+from mev_inspect.schemas.classified_traces import ClassifiedTrace, Classification, NestedTrace
 from mev_inspect.traces import as_nested_traces
 
 
@@ -82,20 +83,13 @@ class TestTraces(unittest.TestCase):
 def build_trace_at_address(
     transaction_hash: str,
     trace_address: List[int],
-) -> Trace:
-    return Trace(
-        # real values
+) -> ClassifiedTrace:
+    return ClassifiedTrace(
         transaction_hash=transaction_hash,
         trace_address=trace_address,
-        # placeholders
-        action={},
-        block_hash="",
         block_number=DEFAULT_BLOCK_NUMBER,
-        result=None,
-        subtraces=0,
-        transaction_position=None,
-        type=TraceType.call,
-        error=None,
+        trace_type=TraceType.call,
+        classification=Classification.unknown,
     )