From d57a2d021d90a172e05f5fab6c2750d2e92999c9 Mon Sep 17 00:00:00 2001 From: Luke Van Seters Date: Wed, 6 Oct 2021 15:12:44 -0400 Subject: [PATCH] Add specific classifiers for each type --- mev_inspect/classifiers/specs/aave.py | 7 ++----- mev_inspect/classifiers/specs/balancer.py | 11 +++-------- mev_inspect/classifiers/specs/erc20.py | 13 +++---------- mev_inspect/classifiers/specs/uniswap.py | 11 +++-------- mev_inspect/classifiers/specs/weth.py | 11 +++-------- mev_inspect/schemas/classifiers.py | 13 +++++++++++++ 6 files changed, 27 insertions(+), 39 deletions(-) diff --git a/mev_inspect/classifiers/specs/aave.py b/mev_inspect/classifiers/specs/aave.py index f8613b1..7a42ed5 100644 --- a/mev_inspect/classifiers/specs/aave.py +++ b/mev_inspect/classifiers/specs/aave.py @@ -1,10 +1,9 @@ from mev_inspect.schemas.classified_traces import ( - Classification, Protocol, ) from mev_inspect.schemas.classifiers import ( ClassifierSpec, - Classifier, + LiquidationClassifier, ) @@ -12,9 +11,7 @@ AAVE_SPEC = ClassifierSpec( abi_name="AaveLendingPool", protocol=Protocol.aave, classifiers={ - "liquidationCall(address,address,address,uint256,bool)": Classifier( - classification=Classification.liquidate, - ) + "liquidationCall(address,address,address,uint256,bool)": LiquidationClassifier(), }, ) diff --git a/mev_inspect/classifiers/specs/balancer.py b/mev_inspect/classifiers/specs/balancer.py index d9c01f9..a9db73c 100644 --- a/mev_inspect/classifiers/specs/balancer.py +++ b/mev_inspect/classifiers/specs/balancer.py @@ -1,10 +1,9 @@ from mev_inspect.schemas.classified_traces import ( - Classification, Protocol, ) from mev_inspect.schemas.classifiers import ( ClassifierSpec, - Classifier, + SwapClassifier, ) @@ -13,12 +12,8 @@ BALANCER_V1_SPECS = [ abi_name="BPool", protocol=Protocol.balancer_v1, classifiers={ - "swapExactAmountIn(address,uint256,address,uint256,uint256)": Classifier( - classification=Classification.swap, - ), - "swapExactAmountOut(address,uint256,address,uint256,uint256)": Classifier( - classification=Classification.swap, - ), + "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 f6c89dc..27093d4 100644 --- a/mev_inspect/classifiers/specs/erc20.py +++ b/mev_inspect/classifiers/specs/erc20.py @@ -1,21 +1,14 @@ -from mev_inspect.schemas.classified_traces import ( - Classification, -) from mev_inspect.schemas.classifiers import ( ClassifierSpec, - Classifier, + TransferClassifier, ) ERC20_SPEC = ClassifierSpec( abi_name="ERC20", classifiers={ - "transferFrom(address,address,uint256)": Classifier( - classification=Classification.transfer, - ), - "transfer(address,uint256)": Classifier( - classification=Classification.transfer, - ), + "transferFrom(address,address,uint256)": TransferClassifier(), + "transfer(address,uint256)": TransferClassifier(), }, ) diff --git a/mev_inspect/classifiers/specs/uniswap.py b/mev_inspect/classifiers/specs/uniswap.py index d02e40d..19b9f16 100644 --- a/mev_inspect/classifiers/specs/uniswap.py +++ b/mev_inspect/classifiers/specs/uniswap.py @@ -1,10 +1,9 @@ from mev_inspect.schemas.classified_traces import ( - Classification, Protocol, ) from mev_inspect.schemas.classifiers import ( ClassifierSpec, - Classifier, + SwapClassifier, ) @@ -70,9 +69,7 @@ UNISWAP_V3_GENERAL_SPECS = [ ClassifierSpec( abi_name="UniswapV3Pool", classifiers={ - "swap(address,bool,int256,uint160,bytes)": Classifier( - classification=Classification.swap, - ) + "swap(address,bool,int256,uint160,bytes)": SwapClassifier(), }, ), ClassifierSpec( @@ -103,9 +100,7 @@ UNISWAPPY_V2_CONTRACT_SPECS = [ UNISWAPPY_V2_PAIR_SPEC = ClassifierSpec( abi_name="UniswapV2Pair", classifiers={ - "swap(uint256,uint256,address,bytes)": Classifier( - classification=Classification.swap, - ), + "swap(uint256,uint256,address,bytes)": SwapClassifier(), }, ) diff --git a/mev_inspect/classifiers/specs/weth.py b/mev_inspect/classifiers/specs/weth.py index f2f68b8..65589d8 100644 --- a/mev_inspect/classifiers/specs/weth.py +++ b/mev_inspect/classifiers/specs/weth.py @@ -1,10 +1,9 @@ from mev_inspect.schemas.classified_traces import ( - Classification, Protocol, ) from mev_inspect.schemas.classifiers import ( ClassifierSpec, - Classifier, + TransferClassifier, ) @@ -13,12 +12,8 @@ WETH_SPEC = ClassifierSpec( protocol=Protocol.weth, valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"], classifiers={ - "transferFrom(address,address,uint256)": Classifier( - classification=Classification.transfer, - ), - "transfer(address,uint256)": Classifier( - classification=Classification.transfer, - ), + "transferFrom(address,address,uint256)": TransferClassifier(), + "transfer(address,uint256)": TransferClassifier(), }, ) diff --git a/mev_inspect/schemas/classifiers.py b/mev_inspect/schemas/classifiers.py index 006ab58..16741f1 100644 --- a/mev_inspect/schemas/classifiers.py +++ b/mev_inspect/schemas/classifiers.py @@ -1,4 +1,5 @@ from typing import Dict, List, Optional +from typing_extensions import Literal from pydantic import BaseModel @@ -9,6 +10,18 @@ class Classifier(BaseModel): classification: Classification +class TransferClassifier(Classifier): + classification: Literal[Classification.transfer] = Classification.transfer + + +class SwapClassifier(Classifier): + classification: Literal[Classification.swap] = Classification.swap + + +class LiquidationClassifier(Classifier): + classification: Literal[Classification.liquidate] = Classification.liquidate + + class ClassifierSpec(BaseModel): abi_name: str protocol: Optional[Protocol] = None