Merge pull request #73 from emlazzarin/main

Balancer v1 classifier / abis

mev_inspect/classifiers/specs/__init__.py

@@ -4,6 +4,7 @@ from .erc20 import ERC20_CLASSIFIER_SPECS
 from .uniswap import UNISWAP_CLASSIFIER_SPECS
 from .weth import WETH_CLASSIFIER_SPECS
 from .zero_ex import ZEROX_CLASSIFIER_SPECS
+from .balancer import BALANCER_CLASSIFIER_SPECS
 
 
 ALL_CLASSIFIER_SPECS = (
@@ -13,4 +14,5 @@ ALL_CLASSIFIER_SPECS = (
     + UNISWAP_CLASSIFIER_SPECS
     + AAVE_CLASSIFIER_SPECS
     + ZEROX_CLASSIFIER_SPECS
+    + BALANCER_CLASSIFIER_SPECS
 )

mev_inspect/classifiers/specs/balancer.py

@@ -0,0 +1,20 @@
+from mev_inspect.schemas.classified_traces import (
+    Classification,
+    ClassifierSpec,
+    Protocol,
+)
+
+BALANCER_V1_SPECS = [
+    ClassifierSpec(
+        abi_name="BPool",
+        protocol=Protocol.balancer_v1,
+        classifications={
+            "swapExactAmountIn(address,uint256,address,uint256,uint256)": Classification.swap,
+            "swapExactAmountOut(address,uint256,address,uint256,uint256)": Classification.swap,
+        },
+    ),
+]
+
+BALANCER_CLASSIFIER_SPECS = [
+    *BALANCER_V1_SPECS,
+]

mev_inspect/schemas/classified_traces.py

@@ -22,6 +22,7 @@ class Protocol(Enum):
     weth = "weth"
     curve = "curve"
     zero_ex = "0x"
+    balancer_v1 = "balancer_v1"
 
 
 class ClassifiedTrace(Trace):

mev_inspect/swaps.py

@@ -16,6 +16,7 @@ from mev_inspect.transfers import (
 
 UNISWAP_V2_PAIR_ABI_NAME = "UniswapV2Pair"
 UNISWAP_V3_POOL_ABI_NAME = "UniswapV3Pool"
+BALANCER_V1_POOL_ABI_NAME = "BPool"
 
 
 def get_swaps(traces: List[ClassifiedTrace]) -> List[Swap]:
@@ -114,5 +115,7 @@ def _get_recipient_address(trace: ClassifiedTrace) -> Optional[str]:
             if trace.inputs is not None and "to" in trace.inputs
             else trace.from_address
         )
+    elif trace.abi_name == BALANCER_V1_POOL_ABI_NAME:
+        return trace.from_address
     else:
         return None

tests/test_swaps.py

@@ -2,6 +2,7 @@ from mev_inspect.swaps import (
     get_swaps,
     UNISWAP_V2_PAIR_ABI_NAME,
     UNISWAP_V3_POOL_ABI_NAME,
+    BALANCER_V1_POOL_ABI_NAME,
 )
 
 from .helpers import (
@@ -20,7 +21,8 @@ def test_swaps(
     [
         first_transaction_hash,
         second_transaction_hash,
-    ] = get_transaction_hashes(2)
+        third_transaction_hash,
+    ] = get_transaction_hashes(3)
 
     [
         alice_address,
@@ -32,12 +34,17 @@ def test_swaps(
         second_token_in_address,
         second_token_out_address,
         second_pool_address,
-    ] = get_addresses(9)
+        third_token_in_address,
+        third_token_out_address,
+        third_pool_address,
+    ] = get_addresses(12)
 
     first_token_in_amount = 10
     first_token_out_amount = 20
     second_token_in_amount = 30
     second_token_out_amount = 40
+    third_token_in_amount = 50
+    third_token_out_amount = 60
 
     traces = [
         make_unknown_trace(block_number, first_transaction_hash, []),
@@ -97,39 +104,85 @@ def test_swaps(
             token_address=second_token_in_address,
             amount=second_token_in_amount,
         ),
+        make_transfer_trace(
+            block_number,
+            third_transaction_hash,
+            trace_address=[6, 0],
+            from_address=bob_address,
+            to_address=third_pool_address,
+            token_address=third_token_in_address,
+            amount=third_token_in_amount,
+        ),
+        make_transfer_trace(
+            block_number,
+            third_transaction_hash,
+            trace_address=[6, 1],
+            from_address=third_pool_address,
+            to_address=bob_address,
+            token_address=third_token_out_address,
+            amount=third_token_out_amount,
+        ),
+        make_swap_trace(
+            block_number,
+            third_transaction_hash,
+            trace_address=[6],
+            from_address=bob_address,
+            pool_address=third_pool_address,
+            abi_name=BALANCER_V1_POOL_ABI_NAME,
+            recipient_address=bob_address,
+            recipient_input_key="recipient",
+        ),
     ]
 
     swaps = get_swaps(traces)
 
-    assert len(swaps) == 2
+    assert len(swaps) == 3
 
-    if swaps[0].abi_name == UNISWAP_V2_PAIR_ABI_NAME:
-        [first_swap, second_swap] = swaps  # pylint: disable=unbalanced-tuple-unpacking
-    else:
-        [second_swap, first_swap] = swaps  # pylint: disable=unbalanced-tuple-unpacking
+    for swap in swaps:
+        if swap.abi_name == UNISWAP_V2_PAIR_ABI_NAME:
+            uni_v2_swap = swap
+        elif swap.abi_name == UNISWAP_V3_POOL_ABI_NAME:
+            uni_v3_swap = swap
+        elif swap.abi_name == BALANCER_V1_POOL_ABI_NAME:
+            bal_v1_swap = swap
+        else:
+            assert False
 
-    assert first_swap.abi_name == UNISWAP_V2_PAIR_ABI_NAME
-    assert first_swap.transaction_hash == first_transaction_hash
-    assert first_swap.block_number == block_number
-    assert first_swap.trace_address == [1]
-    assert first_swap.protocol is None
-    assert first_swap.pool_address == first_pool_address
-    assert first_swap.from_address == alice_address
-    assert first_swap.to_address == bob_address
-    assert first_swap.token_in_address == first_token_in_address
-    assert first_swap.token_in_amount == first_token_in_amount
-    assert first_swap.token_out_address == first_token_out_address
-    assert first_swap.token_out_amount == first_token_out_amount
+    assert uni_v2_swap.abi_name == UNISWAP_V2_PAIR_ABI_NAME
+    assert uni_v2_swap.transaction_hash == first_transaction_hash
+    assert uni_v2_swap.block_number == block_number
+    assert uni_v2_swap.trace_address == [1]
+    assert uni_v2_swap.protocol is None
+    assert uni_v2_swap.pool_address == first_pool_address
+    assert uni_v2_swap.from_address == alice_address
+    assert uni_v2_swap.to_address == bob_address
+    assert uni_v2_swap.token_in_address == first_token_in_address
+    assert uni_v2_swap.token_in_amount == first_token_in_amount
+    assert uni_v2_swap.token_out_address == first_token_out_address
+    assert uni_v2_swap.token_out_amount == first_token_out_amount
 
-    assert second_swap.abi_name == UNISWAP_V3_POOL_ABI_NAME
-    assert second_swap.transaction_hash == second_transaction_hash
-    assert second_swap.block_number == block_number
-    assert second_swap.trace_address == []
-    assert second_swap.protocol is None
-    assert second_swap.pool_address == second_pool_address
-    assert second_swap.from_address == bob_address
-    assert second_swap.to_address == carl_address
-    assert second_swap.token_in_address == second_token_in_address
-    assert second_swap.token_in_amount == second_token_in_amount
-    assert second_swap.token_out_address == second_token_out_address
-    assert second_swap.token_out_amount == second_token_out_amount
+    assert uni_v3_swap.abi_name == UNISWAP_V3_POOL_ABI_NAME
+    assert uni_v3_swap.transaction_hash == second_transaction_hash
+    assert uni_v3_swap.block_number == block_number
+    assert uni_v3_swap.trace_address == []
+    assert uni_v3_swap.protocol is None
+    assert uni_v3_swap.pool_address == second_pool_address
+    assert uni_v3_swap.from_address == bob_address
+    assert uni_v3_swap.to_address == carl_address
+    assert uni_v3_swap.token_in_address == second_token_in_address
+    assert uni_v3_swap.token_in_amount == second_token_in_amount
+    assert uni_v3_swap.token_out_address == second_token_out_address
+    assert uni_v3_swap.token_out_amount == second_token_out_amount
+
+    assert bal_v1_swap.abi_name == BALANCER_V1_POOL_ABI_NAME
+    assert bal_v1_swap.transaction_hash == third_transaction_hash
+    assert bal_v1_swap.block_number == block_number
+    assert bal_v1_swap.trace_address == [6]
+    assert bal_v1_swap.protocol is None
+    assert bal_v1_swap.pool_address == third_pool_address
+    assert bal_v1_swap.from_address == bob_address
+    assert bal_v1_swap.to_address == bob_address
+    assert bal_v1_swap.token_in_address == third_token_in_address
+    assert bal_v1_swap.token_in_amount == third_token_in_amount
+    assert bal_v1_swap.token_out_address == third_token_out_address
+    assert bal_v1_swap.token_out_amount == third_token_out_amount