Fix circular imports

mev_inspect/classifiers/specs/uniswap.py

@@ -1,6 +1,7 @@
-from typing import Optional, List
+from typing import Optional, List, Sequence
 
 from mev_inspect.schemas.traces import (
+    ClassifiedTrace,
     DecodedCallTrace,
     Protocol,
 )
@@ -8,12 +9,9 @@ from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
     SwapClassifier,
 )
+
 from mev_inspect.schemas.swaps import Swap
-from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.schemas.transfers import Transfer, ETH_TOKEN_ADDRESS
-from mev_inspect.transfers import (
-    build_eth_transfer,
-    filter_transfers,
-)
 
 UNISWAP_V2_PAIR_ABI_NAME = "UniswapV2Pair"
 UNISWAP_V3_POOL_ABI_NAME = "UniswapV3Pool"
@@ -39,22 +37,22 @@ class UniswapV3SwapClassifier(SwapClassifier):
         transfers_to_pool = []
 
         if trace.value is not None and trace.value > 0:
-            transfers_to_pool = [build_eth_transfer(trace)]
+            transfers_to_pool = [_build_eth_transfer(trace)]
 
         if len(transfers_to_pool) == 0:
-            transfers_to_pool = filter_transfers(
+            transfers_to_pool = _filter_transfers(
                 prior_transfers, to_address=pool_address
             )
 
         if len(transfers_to_pool) == 0:
-            transfers_to_pool = filter_transfers(
+            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(
+        transfers_from_pool_to_recipient = _filter_transfers(
             child_transfers, to_address=recipient_address, from_address=pool_address
         )
 
@@ -101,22 +99,22 @@ class UniswapV2SwapClassifier(SwapClassifier):
         transfers_to_pool = []
 
         if trace.value is not None and trace.value > 0:
-            transfers_to_pool = [build_eth_transfer(trace)]
+            transfers_to_pool = [_build_eth_transfer(trace)]
 
         if len(transfers_to_pool) == 0:
-            transfers_to_pool = filter_transfers(
+            transfers_to_pool = _filter_transfers(
                 prior_transfers, to_address=pool_address
             )
 
         if len(transfers_to_pool) == 0:
-            transfers_to_pool = filter_transfers(
+            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(
+        transfers_from_pool_to_recipient = _filter_transfers(
             child_transfers, to_address=recipient_address, from_address=pool_address
         )
 
@@ -143,6 +141,37 @@ class UniswapV2SwapClassifier(SwapClassifier):
         )
 
 
+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
+
+
 UNISWAP_V3_CONTRACT_SPECS = [
     ClassifierSpec(
         abi_name="UniswapV3Factory",

mev_inspect/classifiers/specs/zero_ex.py

@@ -5,7 +5,6 @@ from mev_inspect.schemas.classifiers import (
     ClassifierSpec,
 )
 
-
 ZEROX_CONTRACT_SPECS = [
     ClassifierSpec(
         abi_name="exchangeProxy",

mev_inspect/schemas/classifiers.py

@@ -5,6 +5,7 @@ from pydantic import BaseModel
 
 from .traces import Classification, DecodedCallTrace, Protocol
 from .transfers import Transfer
+from .swaps import Swap
 
 
 class Classifier(ABC):
@@ -35,6 +36,15 @@ class SwapClassifier(Classifier):
     def get_swap_recipient(trace: DecodedCallTrace) -> str:
         raise NotImplementedError()
 
+    @staticmethod
+    @abstractmethod
+    def parse_swap(
+        trace: DecodedCallTrace,
+        prior_transfers: List[Transfer],
+        child_transfers: List[Transfer],
+    ) -> Optional[Swap]:
+        raise NotImplementedError()
+
 
 class LiquidationClassifier(Classifier):
     @staticmethod

mev_inspect/swaps.py

@@ -11,10 +11,8 @@ from mev_inspect.schemas.swaps import Swap
 from mev_inspect.schemas.transfers import Transfer
 from mev_inspect.traces import get_traces_by_transaction_hash
 from mev_inspect.transfers import (
-    build_eth_transfer,
     get_child_transfers,
     get_transfer,
-    filter_transfers,
     remove_child_transfers_of_transfers,
 )
 
@@ -67,56 +65,8 @@ def _parse_swap(
     prior_transfers: List[Transfer],
     child_transfers: List[Transfer],
 ) -> Optional[Swap]:
-    pool_address = trace.to_address
-    recipient_address = _get_recipient_address(trace)
 
-    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 _get_recipient_address(trace: DecodedCallTrace) -> Optional[str]:
     classifier = get_classifier(trace)
     if classifier is not None and issubclass(classifier, SwapClassifier):
-        return classifier.get_swap_recipient(trace)
+        return classifier.parse_swap(trace, prior_transfers, child_transfers)
-
     return None