Merge pull request #90 from flashbots/specs-v2

Group classifying a trace as a `transfer` with the logic to decode a `Transfer` object

.gitignore

@@ -19,3 +19,6 @@ cache
 
 # k8s
 .helm
+
+# env
+.envrc

mev_inspect/classifiers/specs/aave.py

@@ -1,14 +1,17 @@
 from mev_inspect.schemas.classified_traces import (
-    Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    LiquidationClassifier,
+)
+
 
 AAVE_SPEC = ClassifierSpec(
     abi_name="AaveLendingPool",
     protocol=Protocol.aave,
-    classifications={
+    classifiers={
-        "liquidationCall(address,address,address,uint256,bool)": Classification.liquidate,
+        "liquidationCall(address,address,address,uint256,bool)": LiquidationClassifier,
     },
 )
 

mev_inspect/classifiers/specs/balancer.py

@@ -1,16 +1,19 @@
 from mev_inspect.schemas.classified_traces import (
-    Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    SwapClassifier,
+)
+
 
 BALANCER_V1_SPECS = [
     ClassifierSpec(
         abi_name="BPool",
         protocol=Protocol.balancer_v1,
-        classifications={
+        classifiers={
-            "swapExactAmountIn(address,uint256,address,uint256,uint256)": Classification.swap,
+            "swapExactAmountIn(address,uint256,address,uint256,uint256)": SwapClassifier,
-            "swapExactAmountOut(address,uint256,address,uint256,uint256)": Classification.swap,
+            "swapExactAmountOut(address,uint256,address,uint256,uint256)": SwapClassifier,
         },
     ),
 ]

mev_inspect/classifiers/specs/curve.py

@@ -1,8 +1,11 @@
 from mev_inspect.schemas.classified_traces import (
-    ClassifierSpec,
     Protocol,
 )
 
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+)
+
 """
 Deployment addresses found here
 https://curve.readthedocs.io/ref-addresses.html

mev_inspect/classifiers/specs/erc20.py

@@ -1,15 +1,30 @@
-from mev_inspect.schemas.classified_traces import (
+from mev_inspect.schemas.classified_traces import DecodedCallTrace
-    Classification,
+from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
+    TransferClassifier,
+)
+from mev_inspect.schemas.transfers import ERC20Transfer
+
+
+class ERC20TransferClassifier(TransferClassifier):
+    @staticmethod
+    def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer:
+        return ERC20Transfer(
+            block_number=trace.block_number,
+            transaction_hash=trace.transaction_hash,
+            trace_address=trace.trace_address,
+            amount=trace.inputs["amount"],
+            to_address=trace.inputs["recipient"],
+            from_address=trace.inputs.get("sender", trace.from_address),
+            token_address=trace.to_address,
         )
 
 
 ERC20_SPEC = ClassifierSpec(
     abi_name="ERC20",
-    classifications={
+    classifiers={
-        "transferFrom(address,address,uint256)": Classification.transfer,
+        "transferFrom(address,address,uint256)": ERC20TransferClassifier,
-        "transfer(address,uint256)": Classification.transfer,
+        "transfer(address,uint256)": ERC20TransferClassifier,
-        "burn(address)": Classification.burn,
     },
 )
 

mev_inspect/classifiers/specs/uniswap.py

@@ -1,8 +1,10 @@
 from mev_inspect.schemas.classified_traces import (
-    Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    SwapClassifier,
+)
 
 
 UNISWAP_V3_CONTRACT_SPECS = [
@@ -66,8 +68,8 @@ UNISWAP_V3_CONTRACT_SPECS = [
 UNISWAP_V3_GENERAL_SPECS = [
     ClassifierSpec(
         abi_name="UniswapV3Pool",
-        classifications={
+        classifiers={
-            "swap(address,bool,int256,uint160,bytes)": Classification.swap,
+            "swap(address,bool,int256,uint160,bytes)": SwapClassifier,
         },
     ),
     ClassifierSpec(
@@ -97,8 +99,8 @@ UNISWAPPY_V2_CONTRACT_SPECS = [
 
 UNISWAPPY_V2_PAIR_SPEC = ClassifierSpec(
     abi_name="UniswapV2Pair",
-    classifications={
+    classifiers={
-        "swap(uint256,uint256,address,bytes)": Classification.swap,
+        "swap(uint256,uint256,address,bytes)": SwapClassifier,
     },
 )
 

mev_inspect/classifiers/specs/weth.py

@@ -1,16 +1,35 @@
 from mev_inspect.schemas.classified_traces import (
-    Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    DecodedCallTrace,
+    TransferClassifier,
+)
+from mev_inspect.schemas.transfers import ERC20Transfer
+
+
+class WethTransferClassifier(TransferClassifier):
+    @staticmethod
+    def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer:
+        return ERC20Transfer(
+            block_number=trace.block_number,
+            transaction_hash=trace.transaction_hash,
+            trace_address=trace.trace_address,
+            amount=trace.inputs["wad"],
+            to_address=trace.inputs["dst"],
+            from_address=trace.from_address,
+            token_address=trace.to_address,
+        )
+
 
 WETH_SPEC = ClassifierSpec(
     abi_name="WETH9",
     protocol=Protocol.weth,
     valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"],
-    classifications={
+    classifiers={
-        "transferFrom(address,address,uint256)": Classification.transfer,
+        "transferFrom(address,address,uint256)": WethTransferClassifier,
-        "transfer(address,uint256)": Classification.transfer,
+        "transfer(address,uint256)": WethTransferClassifier,
     },
 )
 

mev_inspect/classifiers/specs/zero_ex.py

@@ -1,7 +1,9 @@
 from mev_inspect.schemas.classified_traces import (
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+)
 
 
 ZEROX_CONTRACT_SPECS = [

mev_inspect/classifiers/trace.py

@@ -67,8 +67,11 @@ class TraceClassifier:
 
             if call_data is not None:
                 signature = call_data.function_signature
-                classification = spec.classifications.get(
+                classifier = spec.classifiers.get(signature)
-                    signature, Classification.unknown
+                classification = (
+                    Classification.unknown
+                    if classifier is None
+                    else classifier.get_classification()
                 )
 
                 return DecodedCallTrace(

mev_inspect/inspect_block.py

@@ -64,6 +64,8 @@ def inspect_block(
         write_classified_traces(db_session, classified_traces)
 
     transfers = get_transfers(classified_traces)
+    logger.info(f"Found {len(transfers)} transfers")
+
     if should_write_transfers:
         delete_transfers_for_block(db_session, block_number)
         write_transfers(db_session, transfers)

mev_inspect/schemas/classified_traces.py

@@ -1,15 +1,12 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel
-
 from .blocks import Trace
 
 
 class Classification(Enum):
     unknown = "unknown"
     swap = "swap"
-    burn = "burn"
     transfer = "transfer"
     liquidate = "liquidate"
 
@@ -62,12 +59,5 @@ class DecodedCallTrace(CallTrace):
     protocol: Optional[Protocol]
     gas: Optional[int]
     gas_used: Optional[int]
-    function_name: Optional[str]
+    function_name: str
-    function_signature: Optional[str]
+    function_signature: str
-
-
-class ClassifierSpec(BaseModel):
-    abi_name: str
-    protocol: Optional[Protocol] = None
-    valid_contract_addresses: Optional[List[str]] = None
-    classifications: Dict[str, Classification] = {}

mev_inspect/schemas/classifiers.py

@@ -0,0 +1,44 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Type
+
+from pydantic import BaseModel
+
+from .classified_traces import Classification, DecodedCallTrace, Protocol
+from .transfers import ERC20Transfer
+
+
+class Classifier(ABC):
+    @staticmethod
+    @abstractmethod
+    def get_classification() -> Classification:
+        raise NotImplementedError()
+
+
+class TransferClassifier(Classifier):
+    @staticmethod
+    def get_classification() -> Classification:
+        return Classification.transfer
+
+    @staticmethod
+    @abstractmethod
+    def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer:
+        raise NotImplementedError()
+
+
+class SwapClassifier(Classifier):
+    @staticmethod
+    def get_classification() -> Classification:
+        return Classification.swap
+
+
+class LiquidationClassifier(Classifier):
+    @staticmethod
+    def get_classification() -> Classification:
+        return Classification.liquidate
+
+
+class ClassifierSpec(BaseModel):
+    abi_name: str
+    protocol: Optional[Protocol] = None
+    valid_contract_addresses: Optional[List[str]] = None
+    classifiers: Dict[str, Type[Classifier]] = {}

mev_inspect/transfers.py

@@ -1,6 +1,13 @@
-from typing import Dict, List, Optional, Sequence
+from typing import Dict, List, Optional, Sequence, Tuple
 
-from mev_inspect.schemas.classified_traces import Classification, ClassifiedTrace
+from mev_inspect.classifiers.specs import ALL_CLASSIFIER_SPECS
+from mev_inspect.schemas.classifiers import ClassifierSpec, TransferClassifier
+from mev_inspect.schemas.classified_traces import (
+    Classification,
+    ClassifiedTrace,
+    DecodedCallTrace,
+    Protocol,
+)
 from mev_inspect.schemas.transfers import ERC20Transfer, EthTransfer, TransferGeneric
 from mev_inspect.traces import is_child_trace_address, get_child_traces
 
@@ -18,9 +25,21 @@ def get_eth_transfers(traces: List[ClassifiedTrace]) -> List[EthTransfer]:
 def get_transfers(traces: List[ClassifiedTrace]) -> List[ERC20Transfer]:
     transfers = []
 
+    specs_by_abi_name_and_protocol: Dict[
+        Tuple[str, Optional[Protocol]], ClassifierSpec
+    ] = {(spec.abi_name, spec.protocol): spec for spec in ALL_CLASSIFIER_SPECS}
+
     for trace in traces:
-        if trace.classification == Classification.transfer:
+        if not isinstance(trace, DecodedCallTrace):
-            transfers.append(ERC20Transfer.from_trace(trace))
+            continue
+
+        abi_name_and_protocol = (trace.abi_name, trace.protocol)
+        spec = specs_by_abi_name_and_protocol.get(abi_name_and_protocol)
+
+        if spec is not None:
+            classifier = spec.classifiers.get(trace.function_signature)
+            if classifier is not None and issubclass(classifier, TransferClassifier):
+                transfers.append(classifier.get_transfer(trace))
 
     return transfers
 

tests/helpers.py

@@ -56,6 +56,8 @@ def make_swap_trace(
         classification=Classification.swap,
         from_address=from_address,
         to_address=pool_address,
+        function_name="swap",
+        function_signature="swap()",
         inputs={recipient_input_key: recipient_address},
         abi_name=abi_name,
         block_hash=str(block_number),