Move swap logic into classifiers

mev_inspect/classifiers/specs/balancer.py

@@ -1,3 +1,6 @@
+from typing import Optional, List
+from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.schemas.swaps import Swap
 from mev_inspect.schemas.traces import (
     DecodedCallTrace,
     Protocol,
@@ -6,15 +9,21 @@ from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
     SwapClassifier,
 )
-
+from mev_inspect.classifiers.utils import create_swap_from_transfers
 
 BALANCER_V1_POOL_ABI_NAME = "BPool"
 
 
 class BalancerSwapClassifier(SwapClassifier):
     @staticmethod
-    def get_swap_recipient(trace: DecodedCallTrace) -> str:
-        return trace.from_address
+    def parse_swap(
+        trace: DecodedCallTrace,
+        prior_transfers: List[Transfer],
+        child_transfers: List[Transfer],
+    ) -> Optional[Swap]:
+
+        swap = create_swap_from_transfers(trace, prior_transfers, child_transfers)
+        return swap
 
 
 BALANCER_V1_SPECS = [

mev_inspect/classifiers/specs/curve.py

@@ -1,18 +1,28 @@
+from typing import Optional, List
+from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.schemas.swaps import Swap
 from mev_inspect.schemas.traces import (
     Protocol,
+    DecodedCallTrace,
 )
 
 from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
-    DecodedCallTrace,
     SwapClassifier,
 )
+from mev_inspect.classifiers.utils import create_swap_from_transfers
 
 
 class CurveSwapClassifier(SwapClassifier):
     @staticmethod
-    def get_swap_recipient(trace: DecodedCallTrace) -> str:
-        return trace.from_address
+    def parse_swap(
+        trace: DecodedCallTrace,
+        prior_transfers: List[Transfer],
+        child_transfers: List[Transfer],
+    ) -> Optional[Swap]:
+
+        swap = create_swap_from_transfers(trace, prior_transfers, child_transfers)
+        return swap
 
 
 CURVE_BASE_POOLS = [

mev_inspect/classifiers/specs/uniswap.py

@@ -1,7 +1,7 @@
-from typing import Optional, List, Sequence
-
+from typing import Optional, List
+from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.schemas.swaps import Swap
 from mev_inspect.schemas.traces import (
-    ClassifiedTrace,
     DecodedCallTrace,
     Protocol,
 )
@@ -9,9 +9,8 @@ from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
     SwapClassifier,
 )
+from mev_inspect.classifiers.utils import create_swap_from_transfers
 
-from mev_inspect.schemas.swaps import Swap
-from mev_inspect.schemas.transfers import Transfer, ETH_TOKEN_ADDRESS
 
 UNISWAP_V2_PAIR_ABI_NAME = "UniswapV2Pair"
 UNISWAP_V3_POOL_ABI_NAME = "UniswapV3Pool"
@@ -24,59 +23,9 @@ class UniswapV3SwapClassifier(SwapClassifier):
         prior_transfers: List[Transfer],
         child_transfers: List[Transfer],
     ) -> Optional[Swap]:
-        pool_address = trace.to_address
 
-        if trace.inputs is not None and "recipient" in trace.inputs:
-            recipient_address = trace.inputs["recipient"]
-        else:
-            recipient_address = trace.from_address
-
-        if recipient_address is None:
-            return None
-
-        transfers_to_pool = []
-
-        if trace.value is not None and trace.value > 0:
-            transfers_to_pool = [_build_eth_transfer(trace)]
-
-        if len(transfers_to_pool) == 0:
-            transfers_to_pool = _filter_transfers(
-                prior_transfers, to_address=pool_address
-            )
-
-        if len(transfers_to_pool) == 0:
-            transfers_to_pool = _filter_transfers(
-                child_transfers, to_address=pool_address
-            )
-
-        if len(transfers_to_pool) == 0:
-            return None
-
-        transfers_from_pool_to_recipient = _filter_transfers(
-            child_transfers, to_address=recipient_address, from_address=pool_address
-        )
-
-        if len(transfers_from_pool_to_recipient) != 1:
-            return None
-
-        transfer_in = transfers_to_pool[-1]
-        transfer_out = transfers_from_pool_to_recipient[0]
-
-        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,
-            protocol=trace.protocol,
-            from_address=transfer_in.from_address,
-            to_address=transfer_out.to_address,
-            token_in_address=transfer_in.token_address,
-            token_in_amount=transfer_in.amount,
-            token_out_address=transfer_out.token_address,
-            token_out_amount=transfer_out.amount,
-            error=trace.error,
-        )
+        swap = create_swap_from_transfers(trace, prior_transfers, child_transfers)
+        return swap
 
 
 class UniswapV2SwapClassifier(SwapClassifier):
@@ -86,90 +35,9 @@ class UniswapV2SwapClassifier(SwapClassifier):
         prior_transfers: List[Transfer],
         child_transfers: List[Transfer],
     ) -> Optional[Swap]:
-        pool_address = trace.to_address
 
-        if trace.inputs is not None and "to" in trace.inputs:
-            recipient_address = trace.inputs["to"]
-        else:
-            recipient_address = trace.from_address
-
-        if recipient_address is None:
-            return None
-
-        transfers_to_pool = []
-
-        if trace.value is not None and trace.value > 0:
-            transfers_to_pool = [_build_eth_transfer(trace)]
-
-        if len(transfers_to_pool) == 0:
-            transfers_to_pool = _filter_transfers(
-                prior_transfers, to_address=pool_address
-            )
-
-        if len(transfers_to_pool) == 0:
-            transfers_to_pool = _filter_transfers(
-                child_transfers, to_address=pool_address
-            )
-
-        if len(transfers_to_pool) == 0:
-            return None
-
-        transfers_from_pool_to_recipient = _filter_transfers(
-            child_transfers, to_address=recipient_address, from_address=pool_address
-        )
-
-        if len(transfers_from_pool_to_recipient) != 1:
-            return None
-
-        transfer_in = transfers_to_pool[-1]
-        transfer_out = transfers_from_pool_to_recipient[0]
-
-        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,
-            protocol=trace.protocol,
-            from_address=transfer_in.from_address,
-            to_address=transfer_out.to_address,
-            token_in_address=transfer_in.token_address,
-            token_in_amount=transfer_in.amount,
-            token_out_address=transfer_out.token_address,
-            token_out_amount=transfer_out.amount,
-            error=trace.error,
-        )
-
-
-def _build_eth_transfer(trace: ClassifiedTrace) -> Transfer:
-    return Transfer(
-        block_number=trace.block_number,
-        transaction_hash=trace.transaction_hash,
-        trace_address=trace.trace_address,
-        amount=trace.value,
-        to_address=trace.to_address,
-        from_address=trace.from_address,
-        token_address=ETH_TOKEN_ADDRESS,
-    )
-
-
-def _filter_transfers(
-    transfers: Sequence[Transfer],
-    to_address: Optional[str] = None,
-    from_address: Optional[str] = None,
-) -> List[Transfer]:
-    filtered_transfers = []
-
-    for transfer in transfers:
-        if to_address is not None and transfer.to_address != to_address:
-            continue
-
-        if from_address is not None and transfer.from_address != from_address:
-            continue
-
-        filtered_transfers.append(transfer)
-
-    return filtered_transfers
+        swap = create_swap_from_transfers(trace, prior_transfers, child_transfers)
+        return swap
 
 
 UNISWAP_V3_CONTRACT_SPECS = [

mev_inspect/classifiers/utils.py

@@ -0,0 +1,93 @@
+from typing import Optional, List, Sequence
+
+from mev_inspect.schemas.swaps import Swap
+from mev_inspect.schemas.transfers import Transfer, ETH_TOKEN_ADDRESS
+
+from mev_inspect.schemas.traces import DecodedCallTrace, ClassifiedTrace
+
+
+def create_swap_from_transfers(
+    trace: DecodedCallTrace,
+    prior_transfers: List[Transfer],
+    child_transfers: List[Transfer],
+) -> Optional[Swap]:
+    pool_address = trace.to_address
+
+    if trace.inputs is not None and "to" in trace.inputs:
+        recipient_address = trace.inputs["to"]
+    else:
+        recipient_address = trace.from_address
+
+    if recipient_address is None:
+        return None
+
+    transfers_to_pool = []
+
+    if trace.value is not None and trace.value > 0:
+        transfers_to_pool = [_build_eth_transfer(trace)]
+
+    if len(transfers_to_pool) == 0:
+        transfers_to_pool = _filter_transfers(prior_transfers, to_address=pool_address)
+
+    if len(transfers_to_pool) == 0:
+        transfers_to_pool = _filter_transfers(child_transfers, to_address=pool_address)
+
+    if len(transfers_to_pool) == 0:
+        return None
+
+    transfers_from_pool_to_recipient = _filter_transfers(
+        child_transfers, to_address=recipient_address, from_address=pool_address
+    )
+
+    if len(transfers_from_pool_to_recipient) != 1:
+        return None
+
+    transfer_in = transfers_to_pool[-1]
+    transfer_out = transfers_from_pool_to_recipient[0]
+
+    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,
+        protocol=trace.protocol,
+        from_address=transfer_in.from_address,
+        to_address=transfer_out.to_address,
+        token_in_address=transfer_in.token_address,
+        token_in_amount=transfer_in.amount,
+        token_out_address=transfer_out.token_address,
+        token_out_amount=transfer_out.amount,
+        error=trace.error,
+    )
+
+
+def _build_eth_transfer(trace: ClassifiedTrace) -> Transfer:
+    return Transfer(
+        block_number=trace.block_number,
+        transaction_hash=trace.transaction_hash,
+        trace_address=trace.trace_address,
+        amount=trace.value,
+        to_address=trace.to_address,
+        from_address=trace.from_address,
+        token_address=ETH_TOKEN_ADDRESS,
+    )
+
+
+def _filter_transfers(
+    transfers: Sequence[Transfer],
+    to_address: Optional[str] = None,
+    from_address: Optional[str] = None,
+) -> List[Transfer]:
+    filtered_transfers = []
+
+    for transfer in transfers:
+        if to_address is not None and transfer.to_address != to_address:
+            continue
+
+        if from_address is not None and transfer.from_address != from_address:
+            continue
+
+        filtered_transfers.append(transfer)
+
+    return filtered_transfers