diff --git a/mev_inspect/schemas/classifiers.py b/mev_inspect/schemas/classifiers.py index 16741f1..3823603 100644 --- a/mev_inspect/schemas/classifiers.py +++ b/mev_inspect/schemas/classifiers.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Union from typing_extensions import Literal from pydantic import BaseModel @@ -6,22 +6,21 @@ from pydantic import BaseModel from .classified_traces import Classification, Protocol -class Classifier(BaseModel): - classification: Classification - - -class TransferClassifier(Classifier): +class TransferClassifier(BaseModel): classification: Literal[Classification.transfer] = Classification.transfer -class SwapClassifier(Classifier): +class SwapClassifier(BaseModel): classification: Literal[Classification.swap] = Classification.swap -class LiquidationClassifier(Classifier): +class LiquidationClassifier(BaseModel): classification: Literal[Classification.liquidate] = Classification.liquidate +Classifier = Union[TransferClassifier, SwapClassifier, LiquidationClassifier] + + class ClassifierSpec(BaseModel): abi_name: str protocol: Optional[Protocol] = None