From 71b7c99c1754528923e66db8c6422ea3e154e285 Mon Sep 17 00:00:00 2001
From: Gui Heise <gheise@ufl.edu>
Date: Thu, 16 Sep 2021 11:08:16 -0400
Subject: [PATCH] Added pydantic serialization, naming changes

---
 mev_inspect/crud/classified_traces.py    |  4 ++--
 mev_inspect/schemas/classified_traces.py |  6 ------
 mev_inspect/schemas/utils.py             | 19 -------------------
 tests/helpers.py                         | 12 ++++++------
 4 files changed, 8 insertions(+), 33 deletions(-)

diff --git a/mev_inspect/crud/classified_traces.py b/mev_inspect/crud/classified_traces.py
index be0ba49..25c6e2b 100644
--- a/mev_inspect/crud/classified_traces.py
+++ b/mev_inspect/crud/classified_traces.py
@@ -1,6 +1,6 @@
+import json
 from typing import List
 
-from mev_inspect.schemas.utils import to_json
 from mev_inspect.models.classified_traces import ClassifiedTraceModel
 from mev_inspect.schemas.classified_traces import ClassifiedTrace
 
@@ -36,7 +36,7 @@ def write_classified_traces(
                     "abi_name": trace.abi_name,
                     "function_name": trace.function_name,
                     "function_signature": trace.function_signature,
-                    "inputs": to_json(trace.inputs),
+                    "inputs": json.loads(trace.json(include={"inputs"}))["inputs"],
                     "from_address": trace.from_address,
                     "to_address": trace.to_address,
                     "gas": trace.gas,
diff --git a/mev_inspect/schemas/classified_traces.py b/mev_inspect/schemas/classified_traces.py
index 0fcf0ca..c84c32a 100644
--- a/mev_inspect/schemas/classified_traces.py
+++ b/mev_inspect/schemas/classified_traces.py
@@ -25,13 +25,10 @@ class Protocol(Enum):
 
 
 class ClassifiedTrace(Trace):
-
     transaction_hash: str
     block_number: int
     trace_address: List[int]
     classification: Classification
-
-    # Optional
     error: Optional[str]
     to_address: Optional[str]
     from_address: Optional[str]
@@ -54,7 +51,6 @@ class ClassifiedTrace(Trace):
 
 
 class CallTrace(ClassifiedTrace):
-
     to_address: str
     from_address: str
 
@@ -62,8 +58,6 @@ class CallTrace(ClassifiedTrace):
 class DecodedCallTrace(CallTrace):
     inputs: Dict[str, Any]
     abi_name: str
-
-    # Optional
     protocol: Optional[Protocol]
     gas: Optional[int]
     gas_used: Optional[int]
diff --git a/mev_inspect/schemas/utils.py b/mev_inspect/schemas/utils.py
index e3fc2fd..1d15876 100644
--- a/mev_inspect/schemas/utils.py
+++ b/mev_inspect/schemas/utils.py
@@ -15,25 +15,6 @@ def to_original_json_dict(model: BaseModel) -> dict:
     return json.loads(model.json(by_alias=True, exclude_unset=True))
 
 
-def to_json(d):
-    if isinstance(d, dict):
-        tmp = {}
-        for key, value in d.items():
-            if isinstance(value, bytes):
-                tmp[key] = value.hex()
-            if isinstance(value, tuple):
-                res = []
-                for val in value:
-                    if isinstance(val, bytes):
-                        res.append(val.hex())
-                    else:
-                        res.append(val)
-                tmp[key] = res
-
-        d.update(tmp)
-    return json.dumps(d)
-
-
 class Web3Model(BaseModel):
     """BaseModel that handles web3's unserializable objects"""
 
diff --git a/tests/helpers.py b/tests/helpers.py
index c030bbb..72afdef 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -17,8 +17,8 @@ def make_transfer_trace(
     to_address: str,
     token_address: str,
     amount: int,
-    action: dict = {},
-    subtraces: int = 0,
+    action={},
+    subtraces=0,
 ):
     return CallTrace(
         transaction_hash=transaction_hash,
@@ -47,8 +47,8 @@ def make_swap_trace(
     abi_name: str,
     recipient_address: str,
     recipient_input_key: str,
-    action: dict = {},
-    subtraces: int = 0,
+    action={},
+    subtraces=0,
 ):
     return DecodedCallTrace(
         transaction_hash=transaction_hash,
@@ -89,8 +89,8 @@ def make_many_unknown_traces(
     block_number: int,
     transaction_hash: str,
     trace_addresses: List[List[int]],
-    action: dict = {},
-    subtraces: int = 0,
+    action={},
+    subtraces=0,
 ) -> List[ClassifiedTrace]:
 
     return [