Change classifications => classifiers

mev_inspect/classifiers/specs/aave.py

@@ -1,14 +1,20 @@
 from mev_inspect.schemas.classified_traces import (
     Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    Classifier,
+)
+
 
 AAVE_SPEC = ClassifierSpec(
     abi_name="AaveLendingPool",
     protocol=Protocol.aave,
-    classifications={
+    classifiers={
-        "liquidationCall(address,address,address,uint256,bool)": Classification.liquidate,
+        "liquidationCall(address,address,address,uint256,bool)": Classifier(
+            classification=Classification.liquidate,
+        )
     },
 )
 

mev_inspect/classifiers/specs/balancer.py

@@ -1,16 +1,24 @@
 from mev_inspect.schemas.classified_traces import (
     Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    Classifier,
+)
+
 
 BALANCER_V1_SPECS = [
     ClassifierSpec(
         abi_name="BPool",
         protocol=Protocol.balancer_v1,
-        classifications={
+        classifiers={
-            "swapExactAmountIn(address,uint256,address,uint256,uint256)": Classification.swap,
+            "swapExactAmountIn(address,uint256,address,uint256,uint256)": Classifier(
-            "swapExactAmountOut(address,uint256,address,uint256,uint256)": Classification.swap,
+                classification=Classification.swap,
+            ),
+            "swapExactAmountOut(address,uint256,address,uint256,uint256)": Classifier(
+                classification=Classification.swap,
+            ),
         },
     ),
 ]

mev_inspect/classifiers/specs/curve.py

@@ -1,8 +1,11 @@
 from mev_inspect.schemas.classified_traces import (
-    ClassifierSpec,
     Protocol,
 )
 
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+)
+
 """
 Deployment addresses found here
 https://curve.readthedocs.io/ref-addresses.html

mev_inspect/classifiers/specs/erc20.py

@@ -1,15 +1,24 @@
 from mev_inspect.schemas.classified_traces import (
     Classification,
+)
+from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
+    Classifier,
 )
 
 
 ERC20_SPEC = ClassifierSpec(
     abi_name="ERC20",
-    classifications={
+    classifiers={
-        "transferFrom(address,address,uint256)": Classification.transfer,
+        "transferFrom(address,address,uint256)": Classifier(
-        "transfer(address,uint256)": Classification.transfer,
+            classification=Classification.transfer,
-        "burn(address)": Classification.burn,
+        ),
+        "transfer(address,uint256)": Classifier(
+            classification=Classification.transfer,
+        ),
+        "burn(address)": Classifier(
+            classification=Classification.burn,
+        ),
     },
 )
 

mev_inspect/classifiers/specs/uniswap.py

@@ -1,8 +1,11 @@
 from mev_inspect.schemas.classified_traces import (
     Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    Classifier,
+)
 
 
 UNISWAP_V3_CONTRACT_SPECS = [
@@ -66,8 +69,10 @@ UNISWAP_V3_CONTRACT_SPECS = [
 UNISWAP_V3_GENERAL_SPECS = [
     ClassifierSpec(
         abi_name="UniswapV3Pool",
-        classifications={
+        classifiers={
-            "swap(address,bool,int256,uint160,bytes)": Classification.swap,
+            "swap(address,bool,int256,uint160,bytes)": Classifier(
+                classification=Classification.swap,
+            )
         },
     ),
     ClassifierSpec(
@@ -97,8 +102,10 @@ UNISWAPPY_V2_CONTRACT_SPECS = [
 
 UNISWAPPY_V2_PAIR_SPEC = ClassifierSpec(
     abi_name="UniswapV2Pair",
-    classifications={
+    classifiers={
-        "swap(uint256,uint256,address,bytes)": Classification.swap,
+        "swap(uint256,uint256,address,bytes)": Classifier(
+            classification=Classification.swap,
+        ),
     },
 )
 

mev_inspect/classifiers/specs/weth.py

@@ -1,16 +1,24 @@
 from mev_inspect.schemas.classified_traces import (
     Classification,
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+    Classifier,
+)
+
 
 WETH_SPEC = ClassifierSpec(
     abi_name="WETH9",
     protocol=Protocol.weth,
     valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"],
-    classifications={
+    classifiers={
-        "transferFrom(address,address,uint256)": Classification.transfer,
+        "transferFrom(address,address,uint256)": Classifier(
-        "transfer(address,uint256)": Classification.transfer,
+            classification=Classification.transfer,
+        ),
+        "transfer(address,uint256)": Classifier(
+            classification=Classification.transfer,
+        ),
     },
 )
 

mev_inspect/classifiers/specs/zero_ex.py

@@ -1,7 +1,9 @@
 from mev_inspect.schemas.classified_traces import (
-    ClassifierSpec,
     Protocol,
 )
+from mev_inspect.schemas.classifiers import (
+    ClassifierSpec,
+)
 
 
 ZEROX_CONTRACT_SPECS = [

mev_inspect/classifiers/trace.py

@@ -67,8 +67,11 @@ class TraceClassifier:
 
             if call_data is not None:
                 signature = call_data.function_signature
-                classification = spec.classifications.get(
+                classifier = spec.classifiers.get(signature)
-                    signature, Classification.unknown
+                classification = (
+                    Classification.unknown
+                    if classifier is None
+                    else classifier.classification
                 )
 
                 return DecodedCallTrace(

mev_inspect/schemas/classified_traces.py

@@ -1,8 +1,6 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel
-
 from .blocks import Trace
 
 
@@ -64,10 +62,3 @@ class DecodedCallTrace(CallTrace):
     gas_used: Optional[int]
     function_name: Optional[str]
     function_signature: Optional[str]
-
-
-class ClassifierSpec(BaseModel):
-    abi_name: str
-    protocol: Optional[Protocol] = None
-    valid_contract_addresses: Optional[List[str]] = None
-    classifications: Dict[str, Classification] = {}

mev_inspect/schemas/classifiers.py

@@ -0,0 +1,16 @@
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from .classified_traces import Classification, Protocol
+
+
+class Classifier(BaseModel):
+    classification: Classification
+
+
+class ClassifierSpec(BaseModel):
+    abi_name: str
+    protocol: Optional[Protocol] = None
+    valid_contract_addresses: Optional[List[str]] = None
+    classifiers: Dict[str, Classifier] = {}