Switch to class instead of instance

This commit is contained in:
Luke Van Seters 2021-10-06 15:56:28 -04:00
parent 02c9c1cddc
commit f1379cc0a0
10 changed files with 87 additions and 53 deletions

View File

@ -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,
},
)

View File

@ -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,
},
),
]

View File

@ -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,
},
)

View File

@ -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,
},
)

View File

@ -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,
},
)

View File

@ -71,7 +71,7 @@ class TraceClassifier:
classification = (
Classification.unknown
if classifier is None
else classifier.classification
else classifier.get_classification()
)
return DecodedCallTrace(

View File

@ -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)

View File

@ -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

View File

@ -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]] = {}

View File

@ -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