From 588b41333d7a9b3c8eaa0f713fbe58e984984a20 Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Wed, 13 Oct 2021 12:57:14 -0400
Subject: [PATCH] Add Aave transfer classifiers

---
 mev_inspect/aave_liquidations.py      | 21 +++++++--------------
 mev_inspect/classifiers/specs/aave.py | 22 ++++++++++++++++++++--
 2 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/mev_inspect/aave_liquidations.py b/mev_inspect/aave_liquidations.py
index 710aea0..a530ab2 100644
--- a/mev_inspect/aave_liquidations.py
+++ b/mev_inspect/aave_liquidations.py
@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 
 from mev_inspect.traces import (
     get_child_traces,
@@ -11,16 +11,9 @@ from mev_inspect.schemas.classified_traces import (
     Protocol,
 )
 
-<<<<<<< HEAD
-<<<<<<< HEAD
-=======
-from mev_inspect.schemas.transfers import ERC20Transfer, aTokenTransfer, Transfer
->>>>>>> Aave transfers
-=======
-from mev_inspect.schemas.transfers import ERC20Transfer, Transfer
->>>>>>> Aave transfers
-from mev_inspect.schemas.liquidations import Liquidation
 from mev_inspect.transfers import get_transfer
+from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.schemas.liquidations import Liquidation
 
 AAVE_CONTRACT_ADDRESSES: List[str] = [
     # AAVE Proxy
@@ -93,18 +86,18 @@ def _get_payback_token_and_amount(
 
     """Look for and return liquidator payback from liquidation"""
     child: ClassifiedTrace
-    child_transfer: Transfer
-
+    child_transfer: Optional[Transfer]
     for child in child_traces:
 
         if child.classification == Classification.transfer and isinstance(
             child, DecodedCallTrace
         ):
 
-            child_transfer = ERC20Transfer.from_trace(child)
+            child_transfer = get_transfer(child)
 
             if (
-                child_transfer.to_address == liquidator
+                child_transfer is not None
+                and child_transfer.to_address == liquidator
                 and child.from_address in AAVE_CONTRACT_ADDRESSES
             ):
                 return child_transfer.token_address, child_transfer.amount
diff --git a/mev_inspect/classifiers/specs/aave.py b/mev_inspect/classifiers/specs/aave.py
index ef92dc3..f269e96 100644
--- a/mev_inspect/classifiers/specs/aave.py
+++ b/mev_inspect/classifiers/specs/aave.py
@@ -1,10 +1,28 @@
 from mev_inspect.schemas.classified_traces import (
     Protocol,
 )
+
 from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
+    DecodedCallTrace,
+    TransferClassifier,
     LiquidationClassifier,
 )
+from mev_inspect.schemas.transfers import Transfer
+
+
+class AaveTransferClassifier(TransferClassifier):
+    @staticmethod
+    def get_transfer(trace: DecodedCallTrace) -> Transfer:
+        return Transfer(
+            block_number=trace.block_number,
+            transaction_hash=trace.transaction_hash,
+            trace_address=trace.trace_address,
+            amount=trace.inputs["value"],
+            to_address=trace.inputs["to"],
+            from_address=trace.inputs["from"],
+            token_address=trace.to_address,
+        )
 
 
 AAVE_SPEC = ClassifierSpec(
@@ -19,8 +37,8 @@ ATOKENS_SPEC = ClassifierSpec(
     abi_name="aTokens",
     protocol=Protocol.aave,
     classifications={
-        "transferOnLiquidation(address,address,uint256)": Classification.transfer,
-        "transferFrom(address,address,uint256)": Classification.transfer,
+        "transferOnLiquidation(address,address,uint256)": AaveTransferClassifier,
+        "transferFrom(address,address,uint256)": AaveTransferClassifier,
     },
 )