diff --git a/mev_inspect/classifiers/specs/erc20.py b/mev_inspect/classifiers/specs/erc20.py index 27093d4..ae9b8e9 100644 --- a/mev_inspect/classifiers/specs/erc20.py +++ b/mev_inspect/classifiers/specs/erc20.py @@ -1,14 +1,31 @@ +from mev_inspect.schemas.classified_traces import DecodedCallTrace from mev_inspect.schemas.classifiers import ( ClassifierSpec, TransferClassifier, ) +from mev_inspect.schemas.transfers import ERC20Transfer + + +def get_erc20_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_TRANSFER_CLASSIFIER = TransferClassifier(get_transfer=get_erc20_transfer) ERC20_SPEC = ClassifierSpec( abi_name="ERC20", classifiers={ - "transferFrom(address,address,uint256)": TransferClassifier(), - "transfer(address,uint256)": TransferClassifier(), + "transferFrom(address,address,uint256)": ERC20_TRANSFER_CLASSIFIER, + "transfer(address,uint256)": ERC20_TRANSFER_CLASSIFIER, }, ) diff --git a/mev_inspect/classifiers/specs/weth.py b/mev_inspect/classifiers/specs/weth.py index 65589d8..76bd893 100644 --- a/mev_inspect/classifiers/specs/weth.py +++ b/mev_inspect/classifiers/specs/weth.py @@ -3,17 +3,33 @@ from mev_inspect.schemas.classified_traces import ( ) from mev_inspect.schemas.classifiers import ( ClassifierSpec, + DecodedCallTrace, TransferClassifier, ) +from mev_inspect.schemas.transfers import ERC20Transfer +def get_weth_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_TRANSFER_CLASSIFIER = TransferClassifier(get_transfer=get_weth_transfer) + WETH_SPEC = ClassifierSpec( abi_name="WETH9", protocol=Protocol.weth, valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"], classifiers={ - "transferFrom(address,address,uint256)": TransferClassifier(), - "transfer(address,uint256)": TransferClassifier(), + "transferFrom(address,address,uint256)": WETH_TRANSFER_CLASSIFIER, + "transfer(address,uint256)": WETH_TRANSFER_CLASSIFIER, }, ) diff --git a/mev_inspect/schemas/classifiers.py b/mev_inspect/schemas/classifiers.py index 3823603..e3e7a34 100644 --- a/mev_inspect/schemas/classifiers.py +++ b/mev_inspect/schemas/classifiers.py @@ -1,13 +1,15 @@ -from typing import Dict, List, Optional, Union +from typing import Callable, Dict, List, Optional, Union from typing_extensions import Literal from pydantic import BaseModel -from .classified_traces import Classification, Protocol +from .classified_traces import Classification, DecodedCallTrace, Protocol +from .transfers import Transfer class TransferClassifier(BaseModel): classification: Literal[Classification.transfer] = Classification.transfer + get_transfer: Callable[[DecodedCallTrace], Transfer] class SwapClassifier(BaseModel):