From 909ed21a544e0be2458dbee160fc8b013c1c8c2e Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Mon, 23 Aug 2021 10:20:12 -0400
Subject: [PATCH] Add EthTransfer. Make both derive from Transfer

---
 mev_inspect/schemas/transfers.py | 17 ++++++++++++++++-
 mev_inspect/transfers.py         | 26 +++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/mev_inspect/schemas/transfers.py b/mev_inspect/schemas/transfers.py
index c0c0474..af2cf96 100644
--- a/mev_inspect/schemas/transfers.py
+++ b/mev_inspect/schemas/transfers.py
@@ -5,12 +5,27 @@ from pydantic import BaseModel
 from .classified_traces import Classification, ClassifiedTrace, Protocol
 
 
-class ERC20Transfer(BaseModel):
+class Transfer(BaseModel):
     transaction_hash: str
     trace_address: List[int]
     from_address: str
     to_address: str
     amount: int
+
+
+class EthTransfer(Transfer):
+    @classmethod
+    def from_trace(cls, trace: ClassifiedTrace) -> "EthTransfer":
+        return cls(
+            transaction_hash=trace.transaction_hash,
+            trace_address=trace.trace_address,
+            amount=trace.value,
+            to_address=trace.to_address,
+            from_address=trace.from_address,
+        )
+
+
+class ERC20Transfer(Transfer):
     token_address: str
 
     @classmethod
diff --git a/mev_inspect/transfers.py b/mev_inspect/transfers.py
index 0b2eaa6..8476912 100644
--- a/mev_inspect/transfers.py
+++ b/mev_inspect/transfers.py
@@ -1,10 +1,30 @@
 from typing import Dict, List, Optional
 
 from mev_inspect.schemas.classified_traces import Classification, ClassifiedTrace
-from mev_inspect.schemas.transfers import ERC20Transfer
+from mev_inspect.schemas.transfers import ERC20Transfer, EthTransfer, Transfer
 from mev_inspect.traces import is_child_trace_address, get_child_traces
 
 
+def get_eth_transfers(traces: List[ClassifiedTrace]) -> List[EthTransfer]:
+    transfers = []
+
+    for trace in traces:
+        if trace.value is not None and trace.value > 0:
+            transfers.append(EthTransfer.from_trace(trace))
+
+    return transfers
+
+
+def get_transfers(traces: List[ClassifiedTrace]) -> List[ERC20Transfer]:
+    transfers = []
+
+    for trace in traces:
+        if trace.classification == Classification.transfer:
+            transfers.append(ERC20Transfer.from_trace(trace))
+
+    return transfers
+
+
 def get_child_transfers(
     transaction_hash: str,
     parent_trace_address: List[int],
@@ -20,10 +40,10 @@ def get_child_transfers(
 
 
 def filter_transfers(
-    transfers: List[ERC20Transfer],
+    transfers: List[Transfer],
     to_address: Optional[str] = None,
     from_address: Optional[str] = None,
-) -> List[ERC20Transfer]:
+) -> List[Transfer]:
     filtered_transfers = []
 
     for transfer in transfers: