Pass swaps into arbitrage instead of computing inside

mev_inspect/arbitrage.py

@@ -2,35 +2,26 @@ from itertools import groupby
 from typing import List, Optional
 
 from mev_inspect.schemas.arbitrage import Arbitrage
-from mev_inspect.schemas.classified_traces import ClassifiedTrace
 from mev_inspect.schemas.swaps import Swap
-from mev_inspect.swaps import get_swaps
 
 
-def get_arbitrages(traces: List[ClassifiedTrace]) -> List[Arbitrage]:
-    get_transaction_hash = lambda t: t.transaction_hash
-    traces_by_transaction = groupby(
-        sorted(traces, key=get_transaction_hash),
+def get_arbitrages(swaps: List[Swap]) -> List[Arbitrage]:
+    get_transaction_hash = lambda swap: swap.transaction_hash
+    swaps_by_transaction = groupby(
+        sorted(swaps, key=get_transaction_hash),
         key=get_transaction_hash,
     )
 
     all_arbitrages = []
 
-    for _, transaction_traces in traces_by_transaction:
-        all_arbitrages += _get_arbitrages_for_transaction(
-            list(transaction_traces),
+    for _, transaction_swaps in swaps_by_transaction:
+        all_arbitrages += _get_arbitrages_from_swaps(
+            list(transaction_swaps),
         )
 
     return all_arbitrages
 
 
-def _get_arbitrages_for_transaction(
-    traces: List[ClassifiedTrace],
-) -> List[Arbitrage]:
-    swaps = get_swaps(traces)
-    return _get_arbitrages_from_swaps(swaps)
-
-
 def _get_arbitrages_from_swaps(swaps: List[Swap]) -> List[Arbitrage]:
     pool_addresses = {swap.pool_address for swap in swaps}
 

mev_inspect/schemas/swaps.py

@@ -8,6 +8,7 @@ from mev_inspect.schemas.classified_traces import Protocol
 class Swap(BaseModel):
     abi_name: str
     transaction_hash: str
+    block_number: int
     trace_address: List[int]
     protocol: Optional[Protocol]
     pool_address: str

mev_inspect/swaps.py

@@ -1,3 +1,4 @@
+from itertools import groupby
 from typing import List, Optional
 
 from mev_inspect.schemas.classified_traces import (
@@ -18,6 +19,21 @@ UNISWAP_V3_POOL_ABI_NAME = "UniswapV3Pool"
 
 
 def get_swaps(traces: List[ClassifiedTrace]) -> List[Swap]:
+    get_transaction_hash = lambda t: t.transaction_hash
+    traces_by_transaction = groupby(
+        sorted(traces, key=get_transaction_hash),
+        key=get_transaction_hash,
+    )
+
+    swaps = []
+
+    for _, transaction_traces in traces_by_transaction:
+        swaps += _get_swaps_for_transaction(list(transaction_traces))
+
+    return swaps
+
+
+def _get_swaps_for_transaction(traces: List[ClassifiedTrace]) -> List[Swap]:
     ordered_traces = list(sorted(traces, key=lambda t: t.trace_address))
 
     swaps: List[Swap] = []
@@ -73,6 +89,7 @@ def _parse_swap(
     return Swap(
         abi_name=trace.abi_name,
         transaction_hash=trace.transaction_hash,
+        block_number=trace.block_number,
         trace_address=trace.trace_address,
         pool_address=pool_address,
         from_address=transfer_in.from_address,

scripts/inspect_block.py

@@ -12,6 +12,7 @@ from mev_inspect.db import get_session
 from mev_inspect.classifier_specs import CLASSIFIER_SPECS
 from mev_inspect.trace_classifier import TraceClassifier
 from mev_inspect.arbitrage import get_arbitrages
+from mev_inspect.swaps import get_swaps
 
 
 @click.command()
@@ -40,7 +41,10 @@ def inspect_block(block_number: int, rpc: str):
     write_classified_traces(db_session, classified_traces)
     db_session.close()
 
-    arbitrages = get_arbitrages(classified_traces)
+    swaps = get_swaps(classified_traces)
+    print(f"Found {len(swaps)} swaps")
+
+    arbitrages = get_arbitrages(swaps)
     print(f"Found {len(arbitrages)} arbitrages")
 
     stats = get_stats(classified_traces)