From 25e1cbcf8350c1c5d6b1b4bb2450e64f07f4cef0 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Tue, 28 Sep 2021 13:04:46 -0400
Subject: [PATCH] Add transfer classifiers to classifier specs. Implement for
 erc20 and weth

---
 mev_inspect/classifiers/specs/erc20.py  | 21 ++++++++++++++++++++-
 mev_inspect/classifiers/specs/weth.py   | 22 ++++++++++++++++++++++
 mev_inspect/schemas/classifier_specs.py |  7 ++-----
 3 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/mev_inspect/classifiers/specs/erc20.py b/mev_inspect/classifiers/specs/erc20.py
index 297fc53..d5813cb 100644
--- a/mev_inspect/classifiers/specs/erc20.py
+++ b/mev_inspect/classifiers/specs/erc20.py
@@ -1,15 +1,34 @@
+from typing import Optional
+
 from mev_inspect.schemas.classifier_specs import ClassifierSpec
 from mev_inspect.schemas.classified_traces import (
     Classification,
+    DecodedCallTrace,
 )
+from mev_inspect.schemas.transfers import ERC20Transfer
+
+
+def classify_transfer(trace: DecodedCallTrace) -> Optional[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",
+    transfer_classifiers={
+        "transferFrom(address,address,uint256)": classify_transfer,
+        "transfer(address,uint256)": classify_transfer,
+    },
     classifications={
         "transferFrom(address,address,uint256)": Classification.transfer,
         "transfer(address,uint256)": Classification.transfer,
-        "burn(address)": Classification.burn,
     },
 )
 
diff --git a/mev_inspect/classifiers/specs/weth.py b/mev_inspect/classifiers/specs/weth.py
index d9faa0e..21a9074 100644
--- a/mev_inspect/classifiers/specs/weth.py
+++ b/mev_inspect/classifiers/specs/weth.py
@@ -1,17 +1,39 @@
+from typing import Optional
+
 from mev_inspect.schemas.classifier_specs import ClassifierSpec
 from mev_inspect.schemas.classified_traces import (
     Classification,
+    DecodedCallTrace,
     Protocol,
 )
+from mev_inspect.schemas.transfers import ERC20Transfer
+
+
+def classify_transfer(trace: DecodedCallTrace) -> Optional[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_SPEC = ClassifierSpec(
     abi_name="WETH9",
     protocol=Protocol.weth,
     valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"],
+    transfer_classifiers={
+        "transferFrom(address,address,uint256)": classify_transfer,
+        "transfer(address,uint256)": classify_transfer,
+    },
     classifications={
         "transferFrom(address,address,uint256)": Classification.transfer,
         "transfer(address,uint256)": Classification.transfer,
     },
 )
 
+
 WETH_CLASSIFIER_SPECS = [WETH_SPEC]
diff --git a/mev_inspect/schemas/classifier_specs.py b/mev_inspect/schemas/classifier_specs.py
index 7afd941..3e8e4fc 100644
--- a/mev_inspect/schemas/classifier_specs.py
+++ b/mev_inspect/schemas/classifier_specs.py
@@ -3,13 +3,10 @@ from typing import Callable, Dict, List, Optional
 from pydantic import BaseModel
 
 from .classified_traces import Classification, DecodedCallTrace, Protocol
-from .transfers import Transfer
+from .transfers import ERC20Transfer
 
 
-TransferClassifier = Callable[
-    [DecodedCallTrace, List[DecodedCallTrace], List[DecodedCallTrace]],
-    Transfer,
-]
+TransferClassifier = Callable[[DecodedCallTrace], Optional[ERC20Transfer]]
 
 
 class ClassifierSpec(BaseModel):