diff --git a/mev_inspect/classifiers/specs/aave.py b/mev_inspect/classifiers/specs/aave.py index 7a42ed5..4def777 100644 --- a/mev_inspect/classifiers/specs/aave.py +++ b/mev_inspect/classifiers/specs/aave.py @@ -11,7 +11,7 @@ AAVE_SPEC = ClassifierSpec( abi_name="AaveLendingPool", protocol=Protocol.aave, classifiers={ - "liquidationCall(address,address,address,uint256,bool)": LiquidationClassifier(), + "liquidationCall(address,address,address,uint256,bool)": LiquidationClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/balancer.py b/mev_inspect/classifiers/specs/balancer.py index a9db73c..9b05dfd 100644 --- a/mev_inspect/classifiers/specs/balancer.py +++ b/mev_inspect/classifiers/specs/balancer.py @@ -12,8 +12,8 @@ BALANCER_V1_SPECS = [ abi_name="BPool", protocol=Protocol.balancer_v1, classifiers={ - "swapExactAmountIn(address,uint256,address,uint256,uint256)": SwapClassifier(), - "swapExactAmountOut(address,uint256,address,uint256,uint256)": SwapClassifier(), + "swapExactAmountIn(address,uint256,address,uint256,uint256)": SwapClassifier, + "swapExactAmountOut(address,uint256,address,uint256,uint256)": SwapClassifier, }, ), ] diff --git a/mev_inspect/classifiers/specs/erc20.py b/mev_inspect/classifiers/specs/erc20.py index ae9b8e9..aee3179 100644 --- a/mev_inspect/classifiers/specs/erc20.py +++ b/mev_inspect/classifiers/specs/erc20.py @@ -6,26 +6,25 @@ from mev_inspect.schemas.classifiers import ( 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) +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", classifiers={ - "transferFrom(address,address,uint256)": ERC20_TRANSFER_CLASSIFIER, - "transfer(address,uint256)": ERC20_TRANSFER_CLASSIFIER, + "transferFrom(address,address,uint256)": ERC20TransferClassifier, + "transfer(address,uint256)": ERC20TransferClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/uniswap.py b/mev_inspect/classifiers/specs/uniswap.py index 19b9f16..ad6d146 100644 --- a/mev_inspect/classifiers/specs/uniswap.py +++ b/mev_inspect/classifiers/specs/uniswap.py @@ -69,7 +69,7 @@ UNISWAP_V3_GENERAL_SPECS = [ ClassifierSpec( abi_name="UniswapV3Pool", classifiers={ - "swap(address,bool,int256,uint160,bytes)": SwapClassifier(), + "swap(address,bool,int256,uint160,bytes)": SwapClassifier, }, ), ClassifierSpec( @@ -100,7 +100,7 @@ UNISWAPPY_V2_CONTRACT_SPECS = [ UNISWAPPY_V2_PAIR_SPEC = ClassifierSpec( abi_name="UniswapV2Pair", classifiers={ - "swap(uint256,uint256,address,bytes)": SwapClassifier(), + "swap(uint256,uint256,address,bytes)": SwapClassifier, }, ) diff --git a/mev_inspect/classifiers/specs/weth.py b/mev_inspect/classifiers/specs/weth.py index 76bd893..e84f2f3 100644 --- a/mev_inspect/classifiers/specs/weth.py +++ b/mev_inspect/classifiers/specs/weth.py @@ -9,27 +9,27 @@ from mev_inspect.schemas.classifiers import ( 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, - ) +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_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)": WETH_TRANSFER_CLASSIFIER, - "transfer(address,uint256)": WETH_TRANSFER_CLASSIFIER, + "transferFrom(address,address,uint256)": WethTransferClassifier, + "transfer(address,uint256)": WethTransferClassifier, }, ) diff --git a/mev_inspect/classifiers/trace.py b/mev_inspect/classifiers/trace.py index 841ca29..ba7f140 100644 --- a/mev_inspect/classifiers/trace.py +++ b/mev_inspect/classifiers/trace.py @@ -71,7 +71,7 @@ class TraceClassifier: classification = ( Classification.unknown if classifier is None - else classifier.classification + else classifier.get_classification() ) return DecodedCallTrace( diff --git a/mev_inspect/inspect_block.py b/mev_inspect/inspect_block.py index 597e2db..e9af60b 100644 --- a/mev_inspect/inspect_block.py +++ b/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) diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py index 81e3ec2..87336b6 100644 --- a/mev_inspect/schemas/classified_traces.py +++ b/mev_inspect/schemas/classified_traces.py @@ -59,5 +59,5 @@ class DecodedCallTrace(CallTrace): protocol: Optional[Protocol] gas: Optional[int] gas_used: Optional[int] - function_name: Optional[str] - function_signature: Optional[str] + function_name: str + function_signature: str diff --git a/mev_inspect/schemas/classifiers.py b/mev_inspect/schemas/classifiers.py index e3e7a34..ab0b887 100644 --- a/mev_inspect/schemas/classifiers.py +++ b/mev_inspect/schemas/classifiers.py @@ -1,30 +1,44 @@ -from typing import Callable, Dict, List, Optional, Union -from typing_extensions import Literal +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 Transfer +from .transfers import ERC20Transfer -class TransferClassifier(BaseModel): - classification: Literal[Classification.transfer] = Classification.transfer - get_transfer: Callable[[DecodedCallTrace], Transfer] +class Classifier(ABC): + @staticmethod + @abstractmethod + def get_classification() -> Classification: + pass -class SwapClassifier(BaseModel): - classification: Literal[Classification.swap] = Classification.swap +class TransferClassifier(Classifier): + @staticmethod + def get_classification() -> Classification: + return Classification.transfer + + @staticmethod + @abstractmethod + def get_transfer(trace: DecodedCallTrace) -> ERC20Transfer: + pass -class LiquidationClassifier(BaseModel): - classification: Literal[Classification.liquidate] = Classification.liquidate +class SwapClassifier(Classifier): + @staticmethod + def get_classification() -> Classification: + return Classification.swap -Classifier = Union[TransferClassifier, SwapClassifier, LiquidationClassifier] +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, Classifier] = {} + classifiers: Dict[str, Type[Classifier]] = {} diff --git a/mev_inspect/transfers.py b/mev_inspect/transfers.py index fbe3a48..90959cd 100644 --- a/mev_inspect/transfers.py +++ b/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: - transfers.append(ERC20Transfer.from_trace(trace)) + if not isinstance(trace, DecodedCallTrace): + 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