Make Classifier a union

This commit is contained in:
Luke Van Seters 2021-10-06 15:14:24 -04:00
parent d57a2d021d
commit 86ee26dd1a

View File

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