Add transfers and simplify children

mev_inspect/aave_liquidations.py

@@ -1,6 +1,10 @@
 from typing import List
 
-from mev_inspect.traces import get_child_traces, is_child_trace_address
+from mev_inspect.traces import (
+    get_child_traces,
+    is_child_of_any_address,
+    is_child_trace_address,
+)
 from mev_inspect.schemas.classified_traces import (
     ClassifiedTrace,
     DecodedCallTrace,
@@ -8,6 +12,7 @@ from mev_inspect.schemas.classified_traces import (
     Protocol,
 )
 
+from mev_inspect.schemas.transfers import ERC20Transfer
 from mev_inspect.schemas.liquidations import Liquidation
 
 AAVE_CONTRACT_ADDRESSES: List[str] = [
@@ -22,43 +27,24 @@ AAVE_CONTRACT_ADDRESSES: List[str] = [
 ]
 
 
-def is_liquidator_payback(trace: ClassifiedTrace, liquidator: str) -> bool:
+def _get_liquidator_payback(
-    """Finds liquidator payback """
+    child_traces: List[ClassifiedTrace], liquidator: str
-
+) -> int:
-    if isinstance(trace, DecodedCallTrace):
+    for child in child_traces:
-        if "recipient" in trace.inputs:
+        if child.classification == Classification.transfer:
-
+            print("HERE 1")
-            if (
+            child_transfer = ERC20Transfer.from_trace(child)
-                trace.inputs["recipient"] == liquidator
+            print(child_transfer.from_address in AAVE_CONTRACT_ADDRESSES)
-                and trace.from_address in AAVE_CONTRACT_ADDRESSES
+            print(child_transfer.to_address == liquidator)
+            if (child_transfer.to_address == liquidator) and (
+                child.from_address in AAVE_CONTRACT_ADDRESSES
             ):
-                return True
+                return child_transfer.amount
 
-        elif "dst" in trace.inputs:
+    return 0
-            if (
-                trace.inputs["dst"] == liquidator
-                and trace.from_address in AAVE_CONTRACT_ADDRESSES
-            ):
-                return True
-
-    return False
 
 
-def get_received_amount(trace: DecodedCallTrace) -> int:
+def _is_child_of_any_address(
-
-    if "amount" in trace.inputs:
-        received_amount = int(trace.inputs["amount"])
-
-    elif "wad" in trace.inputs:
-        received_amount = int(trace.inputs["wad"])
-
-    else:
-        received_amount = 0
-
-    return received_amount
-
-
-def is_child_of_liquidation(
     trace: ClassifiedTrace, parent_liquidations: List[List[int]]
 ) -> bool:
 
@@ -83,7 +69,7 @@ def get_liquidations(
         if (
             trace.classification == Classification.liquidate
             and isinstance(trace, DecodedCallTrace)
-            and not is_child_of_liquidation(trace, parent_liquidations)
+            and not is_child_of_any_address(trace, parent_liquidations)
         ):
 
             parent_liquidations.append(trace.trace_address)
@@ -93,16 +79,7 @@ def get_liquidations(
                 trace.transaction_hash, trace.trace_address, traces
             )
 
-            for child in child_traces:
+            received_amount = _get_liquidator_payback(child_traces, liquidator)
-
-                if is_liquidator_payback(child, liquidator):
-
-                    assert isinstance(child, DecodedCallTrace)
-                    received_amount = get_received_amount(child)
-
-                else:
-
-                    received_amount = 0
 
             liquidations.append(
                 Liquidation(
@@ -120,4 +97,5 @@ def get_liquidations(
                 )
             )
 
+    print(liquidations)
     return liquidations

mev_inspect/traces.py

@@ -34,6 +34,18 @@ def get_child_traces(
     return child_traces
 
 
+def is_child_of_any_address(
+    trace: ClassifiedTrace, parent_liquidations: List[List[int]]
+) -> bool:
+
+    return any(
+        [
+            is_child_trace_address(trace.trace_address, parent)
+            for parent in parent_liquidations
+        ]
+    )
+
+
 def get_traces_by_transaction_hash(
     traces: List[ClassifiedTrace],
 ) -> Dict[str, List[ClassifiedTrace]]: