From 31254350b4b15d022b43678c787d01a6e47a902a Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Fri, 27 Aug 2021 13:45:39 -0400
Subject: [PATCH] Fancy looks

---
 mev_inspect/schemas/classified_traces.py | 25 +++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py
index 3c6089c..a2f0a37 100644
--- a/mev_inspect/schemas/classified_traces.py
+++ b/mev_inspect/schemas/classified_traces.py
@@ -37,6 +37,8 @@ class ClassifiedTrace(Trace):
     block_number: int
     trace_address: List[int]
     classification: Classification
+
+    # Optional
     error: Optional[str]
     to_address: Optional[str]
     from_address: Optional[str]
@@ -49,13 +51,6 @@ class ClassifiedTrace(Trace):
     inputs: Optional[Dict[str, Any]]
     abi_name: Optional[str]
 
-    class Config:
-        json_encoders = {
-            # a little lazy but fine for now
-            # this is used for bytes value inputs
-            bytes: lambda b: b.hex(),
-        }
-
 
 class Call(ClassifiedTrace):
 
@@ -68,6 +63,8 @@ class ClassifiedCall(ClassifiedTrace):
     inputs: Dict[str, Any]
     abi_name: str
 
+    # Optional
+    protocol: Optional[Protocol]
     gas: Optional[int]
     gas_used: Optional[int]
     function_name: Optional[str]
@@ -78,11 +75,11 @@ class ClassifiedCall(ClassifiedTrace):
 
 
 class Swap(BaseModel):
+
     abi_name: str
     transaction_hash: str
     block_number: int
     trace_address: List[int]
-    protocol: Optional[Protocol]
     pool_address: str
     from_address: str
     to_address: str
@@ -90,10 +87,14 @@ class Swap(BaseModel):
     token_in_amount: int
     token_out_address: str
     token_out_amount: int
+
+    # Optional
+    protocol: Optional[Protocol]
     error: Optional[str]
 
 
 class Arbitrage(BaseModel):
+
     swaps: List[Swap]
     block_number: int
     transaction_hash: str
@@ -104,10 +105,11 @@ class Arbitrage(BaseModel):
     profit_amount: int
 
 
-# -------- Transfers------------------------------------------------------------------------------
+# ------ Transfers ------------------------------------------------------------------------------
 
 
 class Transfer(BaseModel):
+
     transaction_hash: str
     trace_address: List[int]
     from_address: str
@@ -144,17 +146,20 @@ class Transfer(BaseModel):
 
 
 class LiquidationType(Enum):
+
     compound_v2_ceth_liquidation = "compound_v2_ceth_liquidation"
     compound_v2_ctoken_liquidation = "compound_v2_ctoken_liquidation"  # TODO: add logic to handle ctoken liquidations
 
 
 class LiquidationStatus(Enum):
+
     seized = "seized"  # succesfully completed
     check = "check"  # just a liquidation check. i.e searcher only checks if opportunity is still available and reverts accordingly
     out_of_gas = "out_of_gas"  # tx ran out of gas
 
 
 class LiquidationCollateralSource(Enum):
+
     aave_flashloan = "aave_flashloan"
     dydx_flashloan = "dydx_flashloan"
     uniswap_flashloan = "uniswap_flashloan"
@@ -164,6 +169,7 @@ class LiquidationCollateralSource(Enum):
 
 
 class Liquidation(CamelModel):
+
     tx_hash: str
     borrower: str  # account that got liquidated
     collateral_provided: str  # collateral provided by searcher, 'ether' or token contract address
@@ -182,6 +188,7 @@ class Liquidation(CamelModel):
 
 
 class ClassifierSpec(BaseModel):
+
     abi_name: str
     protocol: Optional[Protocol] = None
     valid_contract_addresses: Optional[List[str]] = None