fix(Python): Check for supportsUnbalancedLiquidity in pool state for …

…unbalanced operations.
balancer · Dec 19, 2024 · b0162ed · b0162ed
1 parent a7e8f64
commit b0162ed
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
diff --git a/python/src/add_liquidity.py b/python/src/add_liquidity.py
@@ -4,6 +4,7 @@
     _to_raw_undo_rate_round_up,
     _compute_and_charge_aggregate_swap_fees,
     _get_single_input_index,
+    _require_unbalanced_liquidity_enabled
 )
 from src.base_pool_math import (
     compute_add_liquidity_unbalanced,
@@ -48,6 +49,7 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_
             updated_balances_live_scaled18[i] = a
 
     if add_liquidity_input["kind"] == Kind.UNBALANCED.value:
+        _require_unbalanced_liquidity_enabled(pool_state)
         amounts_in_scaled18 = max_amounts_in_scaled18
         computed = compute_add_liquidity_unbalanced(
             updated_balances_live_scaled18,
@@ -63,6 +65,7 @@ def add_liquidity(add_liquidity_input, pool_state, pool_class, hook_class, hook_
         swap_fee_amounts_scaled18 = computed["swap_fee_amounts"]
 
     elif add_liquidity_input["kind"] == Kind.SINGLE_TOKEN_EXACT_OUT.value:
+        _require_unbalanced_liquidity_enabled(pool_state)
         token_index = _get_single_input_index(max_amounts_in_scaled18)
         amounts_in_scaled18 = max_amounts_in_scaled18
         bpt_amount_out = add_liquidity_input["min_bpt_amount_out_raw"]

diff --git a/python/src/remove_liquidity.py b/python/src/remove_liquidity.py
@@ -4,6 +4,7 @@
     _get_single_input_index,
     _to_raw_undo_rate_round_down,
     _compute_and_charge_aggregate_swap_fees,
+    _require_unbalanced_liquidity_enabled,
 )
 from src.base_pool_math import (
     compute_proportional_amounts_out,
@@ -64,6 +65,7 @@ def remove_liquidity(
             remove_liquidity_input["max_bpt_amount_in_raw"],
         )
     elif remove_liquidity_input["kind"] == RemoveKind.SINGLE_TOKEN_EXACT_IN.value:
+        _require_unbalanced_liquidity_enabled(pool_state)
         bpt_amount_in = remove_liquidity_input["max_bpt_amount_in_raw"]
         amounts_out_scaled18 = min_amounts_out_scaled18
         token_out_index = _get_single_input_index(
@@ -83,6 +85,7 @@ def remove_liquidity(
         amounts_out_scaled18[token_out_index] = computed["amount_out_with_fee"]
         swap_fee_amounts_scaled18 = computed["swap_fee_amounts"]
     elif remove_liquidity_input["kind"] == RemoveKind.SINGLE_TOKEN_EXACT_OUT.value:
+        _require_unbalanced_liquidity_enabled(pool_state)
         amounts_out_scaled18 = min_amounts_out_scaled18
         token_out_index = _get_single_input_index(
             remove_liquidity_input["min_amounts_out_raw"]

diff --git a/python/src/utils.py b/python/src/utils.py
@@ -114,3 +114,8 @@ def _get_single_input_index(max_amounts_in):
         raise ValueError("All zero inputs for single token add")
 
     return input_index
+
+
+def _require_unbalanced_liquidity_enabled(pool_state):
+    if not pool_state.get("supportsUnbalancedLiquidity", True):
+        raise ValueError("DoesNotSupportUnbalancedLiquidity")